From f5e46f621a7be979e5818218cc7a13b7c170006f Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 14 Feb 2026 21:50:12 -0300 Subject: [PATCH 001/147] feat: support per-request enable_json_schema_validation for thread safety Allow passing enable_json_schema_validation as a parameter to completion() and acompletion() instead of only relying on the global litellm.enable_json_schema_validation flag. The per-request value takes priority when provided; otherwise falls back to the global (backward compatible). This makes JSON schema validation safe for concurrent usage in FastAPI and other multi-threaded environments. --- litellm/main.py | 5 + litellm/types/utils.py | 1 + litellm/utils.py | 13 +- .../test_json_schema_validation.py | 139 ++++++++++++++++++ 4 files changed, 157 insertions(+), 1 deletion(-) create mode 100644 tests/litellm/litellm_core_utils/test_json_schema_validation.py diff --git a/litellm/main.py b/litellm/main.py index bca023e65ec..362e5b8b263 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -416,6 +416,8 @@ async def acompletion( # noqa: PLR0915 web_search_options: Optional[OpenAIWebSearchOptions] = None, # Session management shared_session: Optional["ClientSession"] = None, + # Per-request JSON schema validation (overrides litellm.enable_json_schema_validation) + enable_json_schema_validation: Optional[bool] = None, **kwargs, ) -> Union[ModelResponse, CustomStreamWrapper]: """ @@ -560,6 +562,7 @@ async def acompletion( # noqa: PLR0915 "thinking": thinking, "web_search_options": web_search_options, "shared_session": shared_session, + "enable_json_schema_validation": enable_json_schema_validation, } if custom_llm_provider is None: _, custom_llm_provider, _, _ = get_llm_provider( @@ -1045,6 +1048,8 @@ def completion( # type: ignore # noqa: PLR0915 thinking: Optional[AnthropicThinkingParam] = None, # Session management shared_session: Optional["ClientSession"] = None, + # Per-request JSON schema validation (overrides litellm.enable_json_schema_validation) + enable_json_schema_validation: Optional[bool] = None, **kwargs, ) -> Union[ModelResponse, CustomStreamWrapper]: """ diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e1f780ffcc3..72705b82fc2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2939,6 +2939,7 @@ all_litellm_params = ( "shared_session", "search_tool_name", "order", + "enable_json_schema_validation", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) + list(CustomPricingLiteLLMParams.model_fields.keys()) diff --git a/litellm/utils.py b/litellm/utils.py index 6fdd2d88bca..0af844772b9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1318,7 +1318,18 @@ def post_call_processing( ### POST-CALL RULES ### rules_obj.post_call_rules(input=model_response, model=model) ### JSON SCHEMA VALIDATION ### - if litellm.enable_json_schema_validation is True: + # Per-request flag takes priority over global flag + _per_request_validation = ( + optional_params.get("enable_json_schema_validation") + if optional_params is not None + else None + ) + _enable_json_schema_validation = ( + _per_request_validation + if _per_request_validation is not None + else litellm.enable_json_schema_validation + ) + if _enable_json_schema_validation is True: try: if ( optional_params is not None diff --git a/tests/litellm/litellm_core_utils/test_json_schema_validation.py b/tests/litellm/litellm_core_utils/test_json_schema_validation.py new file mode 100644 index 00000000000..859a1238a61 --- /dev/null +++ b/tests/litellm/litellm_core_utils/test_json_schema_validation.py @@ -0,0 +1,139 @@ +""" +Tests for per-request enable_json_schema_validation parameter. + +Ensures the per-request flag overrides the global litellm.enable_json_schema_validation, +making JSON schema validation thread-safe for concurrent usage. + +Related issue: https://github.com/BerriAI/litellm/issues/XXXX +""" + +import json + +import pytest + +import litellm +from litellm.types.utils import ModelResponse +from litellm.utils import Rules, post_call_processing + + +def _make_response(content: dict) -> ModelResponse: + """Create a ModelResponse with the given content as JSON string.""" + response = ModelResponse() + response.choices[0].message.content = json.dumps(content) + return response + + +def _mock_completion(): + """Mock function with __name__ == 'completion' for post_call_processing.""" + pass + + +_mock_completion.__name__ = "completion" + +# Schema that requires 'title' (string) and 'rating' (integer) +STRICT_SCHEMA = { + "type": "json_schema", + "json_schema": { + "name": "MovieReview", + "schema": { + "type": "object", + "properties": { + "title": {"type": "string"}, + "rating": {"type": "integer"}, + }, + "required": ["title", "rating"], + }, + }, +} + +# Response that does NOT match the schema (wrong field names) +INVALID_RESPONSE = _make_response({"name": "test", "age": 25}) + +# Response that matches the schema +VALID_RESPONSE = _make_response({"title": "Inception", "rating": 9}) + + +@pytest.fixture(autouse=True) +def _reset_global_flag(): + """Reset the global flag before and after each test.""" + original = litellm.enable_json_schema_validation + litellm.enable_json_schema_validation = False + yield + litellm.enable_json_schema_validation = original + + +class TestPerRequestJsonSchemaValidation: + """Test that per-request enable_json_schema_validation overrides the global flag.""" + + def test_global_off_no_per_request_skips_validation(self): + """Global OFF + no per-request flag -> no validation (default behavior).""" + litellm.enable_json_schema_validation = False + # Should NOT raise even though response doesn't match schema + post_call_processing( + INVALID_RESPONSE, + "test-model", + {"response_format": STRICT_SCHEMA}, + _mock_completion, + Rules(), + ) + + def test_per_request_on_overrides_global_off(self): + """Global OFF + per-request ON -> validation runs and catches invalid response.""" + litellm.enable_json_schema_validation = False + with pytest.raises(litellm.JSONSchemaValidationError): + post_call_processing( + INVALID_RESPONSE, + "test-model", + { + "response_format": STRICT_SCHEMA, + "enable_json_schema_validation": True, + }, + _mock_completion, + Rules(), + ) + + def test_per_request_off_overrides_global_on(self): + """Global ON + per-request OFF -> validation skipped (per-request wins).""" + litellm.enable_json_schema_validation = True + # Should NOT raise because per-request says False + post_call_processing( + INVALID_RESPONSE, + "test-model", + { + "response_format": STRICT_SCHEMA, + "enable_json_schema_validation": False, + }, + _mock_completion, + Rules(), + ) + + def test_global_on_no_per_request_validates(self): + """Global ON + no per-request flag -> validation runs (backward compatible).""" + litellm.enable_json_schema_validation = True + with pytest.raises(litellm.JSONSchemaValidationError): + post_call_processing( + INVALID_RESPONSE, + "test-model", + {"response_format": STRICT_SCHEMA}, + _mock_completion, + Rules(), + ) + + def test_valid_response_passes_with_per_request_on(self): + """Per-request ON + valid response -> no error raised.""" + post_call_processing( + VALID_RESPONSE, + "test-model", + { + "response_format": STRICT_SCHEMA, + "enable_json_schema_validation": True, + }, + _mock_completion, + Rules(), + ) + + def test_per_request_flag_is_in_all_litellm_params(self): + """Ensure the param is registered so it doesn't leak to provider APIs.""" + from litellm.types.utils import all_litellm_params + + assert "enable_json_schema_validation" in all_litellm_params From fd4dc028bcd65c8be3476fbcdebc66124d54d66e Mon Sep 17 00:00:00 2001 From: Chesars Date: Wed, 18 Feb 2026 16:22:55 -0300 Subject: [PATCH 002/147] fix(anthropic): remove hardcoded reasoning summary in adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The adapter was injecting `summary: "detailed"` into the reasoning config when routing Anthropic thinking requests to OpenAI's Responses API. Per the OpenAI spec, reasoning.summary is opt-in — it should not be added unless the user explicitly requests it. --- .../adapters/handler.py | 12 +----- ...erimental_pass_through_messages_handler.py | 39 ++++++++++++++++++- 2 files changed, 38 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 73e74c228ba..e6a73fdbb59 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -44,8 +44,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: For OpenAI models, Chat Completions typically does not return reasoning text (only token accounting). To return a thinking-like content block in the - Anthropic response format, we route the request through OpenAI's Responses API - and request a reasoning summary. + Anthropic response format, we route the request through OpenAI's Responses API. """ custom_llm_provider = completion_kwargs.get("custom_llm_provider") if custom_llm_provider is None: @@ -80,16 +79,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if isinstance(reasoning_effort, str) and reasoning_effort: completion_kwargs["reasoning_effort"] = { "effort": reasoning_effort, - "summary": "detailed", } - elif isinstance(reasoning_effort, dict): - if ( - "summary" not in reasoning_effort - and "generate_summary" not in reasoning_effort - ): - updated_reasoning_effort = dict(reasoning_effort) - updated_reasoning_effort["summary"] = "detailed" - completion_kwargs["reasoning_effort"] = updated_reasoning_effort @staticmethod def _prepare_completion_kwargs( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 376d14416a3..a4d5ade568b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -177,8 +177,8 @@ def test_openai_model_with_thinking_converts_to_reasoning_effort(): # Verify reasoning_effort is set (converted from thinking) assert "reasoning_effort" in call_kwargs, "reasoning_effort should be passed to completion" - # reasoning_effort is transformed into a dict with effort and summary fields - expected_reasoning_effort = {"effort": "minimal", "summary": "detailed"} + # reasoning_effort is transformed into a dict with effort field only (no hardcoded summary) + expected_reasoning_effort = {"effort": "minimal"} assert call_kwargs["reasoning_effort"] == expected_reasoning_effort, \ f"reasoning_effort should be {expected_reasoning_effort} for budget_tokens=1024, got {call_kwargs.get('reasoning_effort')}" @@ -218,3 +218,38 @@ class TestThinkingParameterTransformation: assert result == {"reasoning_effort": "minimal"} assert "thinking" not in result + + +class TestNoHardcodedReasoningSummary: + """Tests for issue #20998: adapter must not hardcode reasoning summary. + + Per OpenAI spec, reasoning.summary is opt-in. The adapter should not + inject summary='detailed' when the user didn't request it. + """ + + def test_no_summary_added_when_not_requested(self): + """reasoning_effort dict should only contain 'effort', no 'summary'.""" + from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( + LiteLLMMessagesToCompletionTransformationHandler, + ) + + thinking = {"type": "enabled", "budget_tokens": 5000} + completion_kwargs = {"model": "openai/gpt-5.1", "reasoning_effort": "medium"} + LiteLLMMessagesToCompletionTransformationHandler._route_openai_thinking_to_responses_api_if_needed( + completion_kwargs, thinking=thinking + ) + assert completion_kwargs["reasoning_effort"] == {"effort": "medium"} + assert "summary" not in completion_kwargs["reasoning_effort"] + + def test_model_prefixed_with_responses(self): + """Model should be prefixed with 'responses/' for Responses API routing.""" + from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( + LiteLLMMessagesToCompletionTransformationHandler, + ) + + thinking = {"type": "enabled", "budget_tokens": 5000} + completion_kwargs = {"model": "openai/gpt-5.1", "reasoning_effort": "medium"} + LiteLLMMessagesToCompletionTransformationHandler._route_openai_thinking_to_responses_api_if_needed( + completion_kwargs, thinking=thinking + ) + assert completion_kwargs["model"] == "responses/openai/gpt-5.1" From 5596545fb99743b0e8870c5a5b1a30aa88ccd02d Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 15:01:41 -0300 Subject: [PATCH 003/147] fix(gemini): correct streaming finish_reason for tool calls Gemini returns finishReason="STOP" even when tool calls are present, and sends tool_calls and finishReason in separate streaming chunks. The ModelResponseIterator now tracks tool_calls across chunks and correctly maps finish_reason to "tool_calls" per the OpenAI spec. Fixes #21041 --- .../vertex_and_google_ai_studio_gemini.py | 35 +++ ...emini_streaming_tool_call_finish_reason.py | 232 ++++++++++++++++++ 2 files changed, 267 insertions(+) create mode 100644 tests/test_litellm/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py 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 bef83b6d35e..3c06ba94450 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 @@ -2864,6 +2864,7 @@ class ModelResponseIterator: self.logging_obj = logging_obj self.is_function_call = check_is_function_call(logging_obj) self.cumulative_tool_call_index: int = 0 + self.has_seen_tool_calls: bool = False def chunk_parser(self, chunk: dict) -> Optional["ModelResponseStream"]: try: @@ -2902,6 +2903,40 @@ class ModelResponseIterator: cumulative_tool_call_index=self.cumulative_tool_call_index, ) + # Track whether tool_calls have been seen across streaming chunks. + # Gemini sends tool_calls and finishReason in separate chunks, + # so we need to remember if earlier chunks contained tool_calls + # to correctly set finish_reason="tool_calls" per the OpenAI spec. + if not self.has_seen_tool_calls: + for choice in model_response.choices: + if hasattr(choice, "delta") and choice.delta and choice.delta.tool_calls: + self.has_seen_tool_calls = True + break + + # Handle final chunk with finishReason but no content. + # _process_candidates skips candidates without "content", + # so the finish_reason from the final chunk is lost. + if not model_response.choices and _candidates: + from litellm.types.utils import Delta, StreamingChoices + + for candidate in _candidates: + finish_reason_str = candidate.get("finishReason") + if finish_reason_str is not None: + if self.has_seen_tool_calls: + mapped_finish_reason = "tool_calls" + else: + mapped_finish_reason = VertexGeminiConfig._check_finish_reason( + None, finish_reason_str + ) + choice = StreamingChoices( + finish_reason=mapped_finish_reason, + index=candidate.get("index", 0), + delta=Delta(content=None, role=None), + logprobs=None, + enhancements=None, + ) + model_response.choices.append(choice) + 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 diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py b/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py new file mode 100644 index 00000000000..3f8efd47fa3 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py @@ -0,0 +1,232 @@ +""" +Tests for Gemini streaming tool call finish_reason mapping. + +Gemini returns finishReason: "STOP" even when tool calls are present. +Per the OpenAI spec, finish_reason must be "tool_calls" when the model +called a tool. The ModelResponseIterator must track tool_calls across +streaming chunks and correctly set finish_reason on the final chunk. + +Ref: https://github.com/BerriAI/litellm/issues/21041 +""" + +from unittest.mock import MagicMock + +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, +) + + +def _make_logging_obj(**kwargs): + """Create a minimal mock logging object for ModelResponseIterator.""" + logging_obj = MagicMock() + logging_obj.optional_params = kwargs.get("optional_params", {}) + return logging_obj + + +def test_streaming_tool_call_finish_reason_is_tool_calls(): + """ + When Gemini streams tool calls across two chunks: + - Chunk 1: has tool call parts, no finishReason + - Chunk 2: has finishReason="STOP", no content + + The final chunk must have finish_reason="tool_calls" (not "stop"). + """ + logging_obj = _make_logging_obj() + iterator = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + # Chunk 1: tool call with no finishReason + chunk_with_tool_calls = { + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "get_current_weather", + "args": {"location": "Boston, MA"}, + } + } + ], + "role": "model", + }, + "index": 0, + } + ], + } + + # Chunk 2: finishReason="STOP" with no content + chunk_with_finish_reason = { + "candidates": [ + { + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 50, + "candidatesTokenCount": 20, + "totalTokenCount": 70, + }, + } + + # Process chunk 1 + response1 = iterator.chunk_parser(chunk_with_tool_calls) + assert response1 is not None + assert len(response1.choices) == 1 + assert response1.choices[0].delta.tool_calls is not None + assert response1.choices[0].finish_reason == "tool_calls" + assert iterator.has_seen_tool_calls is True + + # Process chunk 2 (final chunk) + response2 = iterator.chunk_parser(chunk_with_finish_reason) + assert response2 is not None + assert len(response2.choices) == 1 + assert response2.choices[0].finish_reason == "tool_calls" + + +def test_streaming_no_tool_calls_finish_reason_is_stop(): + """ + When Gemini streams a regular text response (no tool calls), + the final chunk with finishReason="STOP" should map to "stop". + """ + logging_obj = _make_logging_obj() + iterator = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + # Chunk 1: text content, no finishReason + chunk_with_text = { + "candidates": [ + { + "content": { + "parts": [{"text": "Hello! How can I help?"}], + "role": "model", + }, + "index": 0, + } + ], + } + + # Chunk 2: finishReason="STOP" with no content + chunk_with_finish_reason = { + "candidates": [ + { + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 8, + "totalTokenCount": 18, + }, + } + + # Process chunk 1 + response1 = iterator.chunk_parser(chunk_with_text) + assert response1 is not None + assert len(response1.choices) == 1 + assert iterator.has_seen_tool_calls is False + + # Process chunk 2 + response2 = iterator.chunk_parser(chunk_with_finish_reason) + assert response2 is not None + assert len(response2.choices) == 1 + assert response2.choices[0].finish_reason == "stop" + + +def test_streaming_multiple_tool_calls_finish_reason(): + """ + When Gemini streams multiple tool calls across chunks, + the final finish_reason must still be "tool_calls". + """ + logging_obj = _make_logging_obj() + iterator = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + # Chunk 1: first tool call + chunk_tool_1 = { + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "get_weather", + "args": {"location": "NYC"}, + } + }, + { + "functionCall": { + "name": "get_time", + "args": {"timezone": "EST"}, + } + }, + ], + "role": "model", + }, + "index": 0, + } + ], + } + + # Chunk 2: finishReason="STOP" with no content + chunk_finish = { + "candidates": [ + { + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 50, + "candidatesTokenCount": 30, + "totalTokenCount": 80, + }, + } + + response1 = iterator.chunk_parser(chunk_tool_1) + assert response1 is not None + assert iterator.has_seen_tool_calls is True + + response2 = iterator.chunk_parser(chunk_finish) + assert response2 is not None + assert len(response2.choices) == 1 + assert response2.choices[0].finish_reason == "tool_calls" + + +def test_streaming_content_filter_finish_reason_preserved(): + """ + When Gemini returns finishReason due to content filtering (not STOP), + and no tool calls were seen, the content_filter reason should be preserved. + """ + logging_obj = _make_logging_obj() + iterator = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj, + ) + + # Chunk with finishReason="SAFETY" and no content + chunk_safety = { + "candidates": [ + { + "finishReason": "SAFETY", + "index": 0, + } + ], + } + + response = iterator.chunk_parser(chunk_safety) + assert response is not None + assert len(response.choices) == 1 + assert response.choices[0].finish_reason == "content_filter" From 34d09a314ca5a60d7cd64a7f1e531213c01c5d90 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Mon, 23 Feb 2026 14:47:07 -0800 Subject: [PATCH 004/147] virtual-keys-team-table --- .../src/components/team/TeamInfo.test.tsx | 59 +- .../src/components/team/TeamInfo.tsx | 12 + .../team/TeamVirtualKeysTable.test.tsx | 192 +++++ .../components/team/TeamVirtualKeysTable.tsx | 757 ++++++++++++++++++ .../team/tabVisibilityUtils.test.ts | 19 +- .../src/components/team/tabVisibilityUtils.ts | 6 +- 6 files changed, 1040 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx create mode 100644 ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index cc0483aafd6..75fe76902bd 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -100,12 +100,28 @@ vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ }), })); +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ + useKeys: vi.fn().mockReturnValue({ + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + }), +})); + +vi.mock("../key_team_helpers/filter_helpers", () => ({ + fetchAllKeyAliases: vi.fn().mockResolvedValue([]), + fetchAllOrganizations: vi.fn().mockResolvedValue([]), +})); + import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; const mockUseAllProxyModels = vi.mocked(useAllProxyModels); +const mockUseKeys = vi.mocked(useKeys); const mockUseTeam = vi.mocked(useTeam); const mockUseOrganization = vi.mocked(useOrganization); const mockUseCurrentUser = vi.mocked(useCurrentUser); @@ -180,6 +196,12 @@ describe("TeamInfoView", () => { data: { models: [] }, isLoading: false, } as any); + mockUseKeys.mockReturnValue({ + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); vi.mocked(networking.getPoliciesList).mockResolvedValue({ policies: [] }); @@ -558,7 +580,42 @@ describe("TeamInfoView", () => { renderWithProviders(); await waitFor(() => { - expect(screen.getByText("Virtual Keys")).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Virtual Keys" })).toBeInTheDocument(); + }); + }); + + it("should show Virtual Keys tab when user cannot edit team", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("tab", { name: "Virtual Keys" })).toBeInTheDocument(); + }); + }); + + it("should display X Members in Virtual Keys tab when navigated to", async () => { + const user = userEvent.setup(); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + mockUseKeys.mockReturnValue({ + data: { keys: [], total_count: 5, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const virtualKeysTab = screen.getByRole("tab", { name: "Virtual Keys" }); + await user.click(virtualKeysTab); + + await waitFor(() => { + expect(screen.getByText("5 Members")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index ea7a9a1c460..62208b4186d 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -48,6 +48,7 @@ import { TEAM_INFO_TAB_LABELS, } from "./tabVisibilityUtils"; import TeamMembersComponent from "./TeamMemberTab"; +import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; export interface TeamMembership { user_id: string; @@ -726,6 +727,17 @@ const TeamInfoView: React.FC = ({ ), }, + { + key: TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS, + label: TEAM_INFO_TAB_LABELS[TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS], + children: ( + + ), + }, { key: TEAM_INFO_TAB_KEYS.MEMBERS, label: TEAM_INFO_TAB_LABELS[TEAM_INFO_TAB_KEYS.MEMBERS], diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx new file mode 100644 index 00000000000..e4b7b00fb58 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -0,0 +1,192 @@ +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi, MockedFunction } from "vitest"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; +import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { KeyResponse } from "../key_team_helpers/key_list"; +import { Organization } from "../networking"; + +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ + useKeys: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(), +})); + +vi.mock("../key_team_helpers/filter_helpers", () => ({ + fetchAllKeyAliases: vi.fn().mockResolvedValue([]), + fetchAllOrganizations: vi.fn().mockResolvedValue([]), +})); + +vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({ + getModelDisplayName: vi.fn((model: string) => model), +})); + +vi.mock("../templates/key_info_view", () => ({ + default: vi.fn(({ onClose }: { onClose: () => void }) => ( +
+ Key Info View + +
+ )), +})); + +const mockUseKeys = useKeys as MockedFunction; +const mockUseAuthorized = useAuthorized as MockedFunction; + +const createMockKey = (overrides: Partial = {}): KeyResponse => + ({ + token: "sk-test123", + token_id: "key-1", + key_alias: "alice_key_team1", + key_name: "sk-...abc", + user_id: "user-1", + organization_id: null, + user: { user_id: "user-1", user_email: "alice@example.com" }, + created_at: "2024-01-01T00:00:00Z", + team_id: "team-1", + spend: 0, + max_budget: 100, + models: ["gpt-4"], + ...overrides, + } as KeyResponse); + +const mockOrganization: Organization = { + organization_id: "org-123", + organization_alias: "Test Org", + budget_id: "budget-1", + metadata: {}, + models: [], + spend: 0, + model_spend: {}, + created_at: "", + created_by: "", + updated_at: "", + updated_by: "", + litellm_budget_table: {}, + teams: [], + users: [], + members: [], +}; + +describe("TeamVirtualKeysTable", () => { + const defaultProps = { + teamId: "team-1", + teamAlias: "team1", + organization: null as Organization | null, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" } as any); + mockUseKeys.mockReturnValue({ + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + }); + + it("should render successfully", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("0 Members")).toBeInTheDocument(); + }); + }); + + it("should display X Members instead of Showing X of Y results", async () => { + mockUseKeys.mockReturnValue({ + data: { + keys: [createMockKey(), createMockKey({ token: "sk-2", token_id: "key-2" })], + total_count: 2, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("2 Members")).toBeInTheDocument(); + }); + expect(screen.queryByText(/Showing.*results/)).not.toBeInTheDocument(); + }); + + it("should display 1 Member when singular", async () => { + mockUseKeys.mockReturnValue({ + data: { + keys: [createMockKey()], + total_count: 1, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("1 Member")).toBeInTheDocument(); + }); + }); + + it("should call useKeys with expand user to fetch user email", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenCalledWith( + 1, + 50, + expect.objectContaining({ + teamID: "team-1", + expand: "user", + }) + ); + }); + }); + + it("should enrich keys with organization_id when organization is provided", async () => { + const keyWithoutOrg = createMockKey({ organization_id: null }); + mockUseKeys.mockReturnValue({ + data: { + keys: [keyWithoutOrg], + total_count: 1, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders( + + ); + + await waitFor(() => { + expect(screen.getByText("1 Member")).toBeInTheDocument(); + }); + // Key with org_id should display in table - org-123 from organization + await waitFor(() => { + expect(screen.getByText("org-123")).toBeInTheDocument(); + }); + }); + + it("should show table with Key ID column header", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("0 Members")).toBeInTheDocument(); + }); + expect(screen.getByText("Key ID")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx new file mode 100644 index 00000000000..ce998d18677 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -0,0 +1,757 @@ +// TO-DO: Standardize tables eventually + +"use client"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; +import { + ColumnDef, + flexRender, + getCoreRowModel, + getPaginationRowModel, + getSortedRowModel, + PaginationState, + SortingState, + useReactTable, +} from "@tanstack/react-table"; +import { + Badge, + Button, + Icon, + Table, + TableBody, + TableCell, + TableHead, + TableHeaderCell, + TableRow, + Text, +} from "@tremor/react"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { Popover, Skeleton, Tooltip } from "antd"; +import React, { useEffect, useMemo, useState } from "react"; +import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; +import FilterComponent, { FilterOption } from "../molecules/filter"; +import { Organization } from "../networking"; +import KeyInfoView from "../templates/key_info_view"; +import { useQuery } from "@tanstack/react-query"; +import { fetchAllKeyAliases, fetchAllOrganizations } from "../key_team_helpers/filter_helpers"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +interface TeamVirtualKeysTableProps { + teamId: string; + teamAlias?: string; + organization: Organization | null; +} + +/** + * TeamVirtualKeysTable – variant of VirtualKeysTable scoped to a single team. + * Displays all virtual keys belonging to the team with same format and styling. + */ +export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVirtualKeysTableProps) { + const { accessToken } = useAuthorized(); + const [selectedKey, setSelectedKey] = useState(null); + const [sorting, setSorting] = useState([ + { id: "created_at", desc: true }, + ]); + const [tablePagination, setTablePagination] = useState({ + pageIndex: 0, + pageSize: 50, + }); + const [filters, setFilters] = useState>({ + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }); + + const sortBy = sorting.length > 0 ? sorting[0].id : "created_at"; + const sortOrder = sorting.length > 0 ? (sorting[0].desc ? "desc" : "asc") : "desc"; + + const { + data: keys, + isPending: isLoading, + isFetching, + refetch, + } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, { + teamID: teamId, + sortBy: sortBy || undefined, + sortOrder: sortOrder || undefined, + organizationID: filters["Organization ID"] || undefined, + selectedKeyAlias: filters["Key Alias"] || undefined, + userID: filters["User ID"] || undefined, + expand: "user", + }); + + const totalCount = keys?.total_count || 0; + const displayKeys = useMemo(() => { + const kList = keys?.keys || []; + const orgId = organization?.organization_id; + if (!orgId) return kList; + return kList.map((k: KeyResponse) => ({ + ...k, + organization_id: k.organization_id || orgId, + })); + }, [keys?.keys, organization?.organization_id]); + const [expandedAccordions, setExpandedAccordions] = useState>({}); + + const currentTeam: Team = useMemo( + () => ({ + team_id: teamId, + team_alias: teamAlias || teamId, + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: organization?.organization_id || "", + created_at: "", + keys: [], + members_with_roles: [], + spend: 0, + }), + [teamId, teamAlias, organization], + ); + + const allKeyAliasesQuery = useQuery({ + queryKey: ["allKeyAliases"], + queryFn: async () => fetchAllKeyAliases(accessToken), + enabled: !!accessToken, + }); + const allKeyAliases = allKeyAliasesQuery.data || []; + + const allOrganizationsQuery = useQuery({ + queryKey: ["allOrganizations"], + queryFn: async () => fetchAllOrganizations(accessToken), + enabled: !!accessToken, + }); + const allOrganizations = allOrganizationsQuery.data || []; + + useEffect(() => { + if (refetch) { + const handleStorageChange = () => refetch(); + window.addEventListener("storage", handleStorageChange); + return () => window.removeEventListener("storage", handleStorageChange); + } + }, [refetch]); + + const handleFilterChange = (newFilters: Record, skipDebounce = false) => { + setFilters((prev) => ({ + ...prev, + "Organization ID": newFilters["Organization ID"] ?? prev["Organization ID"], + "Key Alias": newFilters["Key Alias"] ?? prev["Key Alias"], + "User ID": newFilters["User ID"] ?? prev["User ID"], + "Sort By": newFilters["Sort By"] ?? prev["Sort By"] ?? "created_at", + "Sort Order": newFilters["Sort Order"] ?? prev["Sort Order"] ?? "desc", + })); + if (!skipDebounce) { + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + } + }; + + const handleFilterReset = () => { + setFilters({ + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }); + setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); + }; + + const filterOptions: FilterOption[] = useMemo( + () => [ + { + name: "Organization ID", + label: "Organization ID", + isSearchable: true, + searchFn: async (searchText: string) => { + if (!allOrganizations.length) return []; + const filtered = allOrganizations.filter( + (org) => org.organization_id?.toLowerCase().includes(searchText.toLowerCase()) ?? false, + ); + return filtered + .filter((org) => org.organization_id != null) + .map((org) => ({ + label: `${org.organization_id || "Unknown"} (${org.organization_id})`, + value: org.organization_id as string, + })); + }, + }, + { + name: "Key Alias", + label: "Key Alias", + isSearchable: true, + searchFn: async (searchText: string) => { + const filtered = allKeyAliases.filter((alias) => + alias.toLowerCase().includes(searchText.toLowerCase()), + ); + return filtered.map((alias) => ({ label: alias, value: alias })); + }, + }, + { name: "User ID", label: "User ID", isSearchable: false }, + ], + [allOrganizations, allKeyAliases], + ); + + const columns: ColumnDef[] = useMemo( + () => [ + { + id: "token", + accessorKey: "token", + header: "Key ID", + size: 100, + enableSorting: true, + cell: (info) => { + const value = info.getValue() as string; + const width = info.cell.column.getSize(); + return ( + + + + ); + }, + }, + { + id: "key_alias", + accessorKey: "key_alias", + header: "Key Alias", + size: 150, + enableSorting: true, + cell: (info) => { + const value = info.getValue() as string; + const width = info.cell.column.getSize(); + return ( + + + {value ?? "-"} + + + ); + }, + }, + { + id: "key_name", + accessorKey: "key_name", + header: "Secret Key", + size: 120, + enableSorting: false, + cell: (info) => {info.getValue() as string}, + }, + { + id: "organization_id", + accessorKey: "organization_id", + header: "Organization ID", + size: 140, + enableSorting: false, + cell: (info) => (info.getValue() ? info.renderValue() : "-"), + }, + { + id: "user_email", + accessorKey: "user", + header: "User Email", + size: 160, + enableSorting: false, + cell: (info) => { + const user = info.getValue() as { user_email?: string } | undefined; + const value = user?.user_email; + const width = info.cell.column.getSize(); + return ( + + + {value ?? "-"} + + + ); + }, + }, + { + id: "user_id", + accessorKey: "user_id", + header: "User ID", + size: 70, + enableSorting: false, + cell: (info) => { + const userId = info.getValue() as string | null; + const displayValue = userId === "default_user_id" ? "Default Proxy Admin" : userId; + const width = info.cell.column.getSize(); + return ( + + + {displayValue ?? "-"} + + + ); + }, + }, + { + id: "created_at", + accessorKey: "created_at", + header: "Created At", + size: 120, + enableSorting: true, + cell: (info) => { + const value = info.getValue(); + return value ? new Date(value as string).toLocaleDateString() : "-"; + }, + }, + { + id: "created_by", + accessorKey: "created_by", + header: "Created By", + size: 70, + enableSorting: false, + cell: (info) => { + const value = info.getValue() as string | null; + const displayValue = value === "default_user_id" ? "Default Proxy Admin" : value; + const width = info.cell.column.getSize(); + return ( + + + {displayValue ?? "-"} + + + ); + }, + }, + { + id: "updated_at", + accessorKey: "updated_at", + header: "Updated At", + size: 120, + enableSorting: true, + cell: (info) => { + const value = info.getValue(); + return value ? new Date(value as string).toLocaleDateString() : "Never"; + }, + }, + { + id: "last_active", + accessorKey: "last_active", + header: () => ( + + Last Active + + + + + ), + size: 130, + enableSorting: false, + cell: (info) => { + const value = info.getValue(); + if (!value) return "Unknown"; + const date = new Date(value as string); + return ( + + {date.toLocaleDateString()} + + ); + }, + }, + { + id: "expires", + accessorKey: "expires", + header: "Expires", + size: 120, + enableSorting: false, + cell: (info) => { + const value = info.getValue(); + return value ? new Date(value as string).toLocaleDateString() : "Never"; + }, + }, + { + id: "spend", + accessorKey: "spend", + header: "Spend (USD)", + size: 100, + enableSorting: true, + cell: (info) => formatNumberWithCommas(info.getValue() as number, 4), + }, + { + id: "max_budget", + accessorKey: "max_budget", + header: "Budget (USD)", + size: 110, + enableSorting: true, + cell: (info) => { + const maxBudget = info.getValue() as number | null; + if (maxBudget === null) return "Unlimited"; + return `$${formatNumberWithCommas(maxBudget)}`; + }, + }, + { + id: "budget_reset_at", + accessorKey: "budget_reset_at", + header: "Budget Reset", + size: 130, + enableSorting: false, + cell: (info) => { + const value = info.getValue(); + return value ? new Date(value as string).toLocaleString() : "Never"; + }, + }, + { + id: "models", + accessorKey: "models", + header: "Models", + size: 200, + enableSorting: false, + cell: (info) => { + const models = info.getValue() as string[]; + return ( +
+ {Array.isArray(models) ? ( +
+ {models.length === 0 ? ( + + All Proxy Models + + ) : ( + <> +
+ {models.length > 3 && ( +
+ + setExpandedAccordions((prev) => ({ + ...prev, + [info.row.id]: !prev[info.row.id], + })) + } + /> +
+ )} +
+ {models.slice(0, 3).map((model, index) => + model === "all-proxy-models" ? ( + + All Proxy Models + + ) : ( + + + {model.length > 30 + ? `${getModelDisplayName(model).slice(0, 30)}...` + : getModelDisplayName(model)} + + + ), + )} + {models.length > 3 && !expandedAccordions[info.row.id] && ( + + + +{models.length - 3} {models.length - 3 === 1 ? "more model" : "more models"} + + + )} + {expandedAccordions[info.row.id] && ( +
+ {models.slice(3).map((model, index) => + model === "all-proxy-models" ? ( + + All Proxy Models + + ) : ( + + + {model.length > 30 + ? `${getModelDisplayName(model).slice(0, 30)}...` + : getModelDisplayName(model)} + + + ), + )} +
+ )} +
+
+ + )} +
+ ) : null} +
+ ); + }, + }, + { + id: "rate_limits", + header: "Rate Limits", + size: 140, + enableSorting: false, + cell: ({ row }) => { + const key = row.original; + return ( +
+
TPM: {key.tpm_limit !== null ? key.tpm_limit : "Unlimited"}
+
RPM: {key.rpm_limit !== null ? key.rpm_limit : "Unlimited"}
+
+ ); + }, + }, + ], + [expandedAccordions], + ); + + const table = useReactTable({ + data: displayKeys, + columns, + columnResizeMode: "onChange", + columnResizeDirection: "ltr", + state: { sorting, pagination: tablePagination }, + onSortingChange: (updaterOrValue) => { + const newSorting = + typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue; + setSorting(newSorting); + if (newSorting?.length > 0) { + const sortState = newSorting[0]; + handleFilterChange( + { + ...filters, + "Sort By": sortState.id, + "Sort Order": sortState.desc ? "desc" : "asc", + }, + true, + ); + } + }, + onPaginationChange: setTablePagination, + getCoreRowModel: getCoreRowModel(), + getSortedRowModel: getSortedRowModel(), + getPaginationRowModel: getPaginationRowModel(), + enableSorting: true, + manualSorting: false, + manualPagination: true, + pageCount: Math.ceil(totalCount / tablePagination.pageSize), + }); + + const { pageIndex, pageSize } = table.getState().pagination; + return ( +
+ {selectedKey ? ( + setSelectedKey(null)} + keyData={selectedKey} + teams={[currentTeam]} + onDelete={refetch} + /> + ) : ( +
+
+ +
+ +
+ {isLoading || isFetching ? ( + + ) : ( + + {totalCount} Member{totalCount !== 1 ? "s" : ""} + + )} + +
+ {isLoading || isFetching ? ( + + ) : ( + + Page {pageIndex + 1} of {table.getPageCount()} + + )} + + {isLoading || isFetching ? ( + + ) : ( + + )} + + {isLoading || isFetching ? ( + + ) : ( + + )} +
+
+
+
+
+ + + {table.getHeaderGroups().map((headerGroup) => ( + + {headerGroup.headers.map((header) => ( + { + const resizer = document.querySelector( + `[data-header-id="${header.id}"] .resizer`, + ); + if (resizer) (resizer as HTMLElement).style.opacity = "0.5"; + }} + onMouseLeave={() => { + const resizer = document.querySelector( + `[data-header-id="${header.id}"] .resizer`, + ); + if (resizer && !header.column.getIsResizing()) + (resizer as HTMLElement).style.opacity = "0"; + }} + onClick={ + header.column.getCanSort() + ? header.column.getToggleSortingHandler() + : undefined + } + > +
+
+ {header.isPlaceholder + ? null + : flexRender(header.column.columnDef.header, header.getContext())} +
+ {header.id !== "actions" && header.column.getCanSort() && ( +
+ {header.column.getIsSorted() ? ( + { + asc: , + desc: , + }[header.column.getIsSorted() as string] + ) : ( + + )} +
+ )} +
header.column.resetSize()} + onMouseDown={header.getResizeHandler()} + onTouchStart={header.getResizeHandler()} + className={`resizer ${table.options.columnResizeDirection} ${ + header.column.getIsResizing() ? "isResizing" : "" + }`} + style={{ + position: "absolute", + right: 0, + top: 0, + height: "100%", + width: "5px", + background: header.column.getIsResizing() ? "#3b82f6" : "transparent", + cursor: "col-resize", + userSelect: "none", + touchAction: "none", + opacity: header.column.getIsResizing() ? 1 : 0, + }} + /> +
+ + ))} + + ))} + + + {isLoading || isFetching ? ( + + +
+

Loading keys...

+
+ + + ) : displayKeys.length > 0 ? ( + table.getRowModel().rows.map((row) => ( + + {row.getVisibleCells().map((cell) => ( + 3 + ? "px-0" + : "" + }`} + > + {flexRender(cell.column.columnDef.cell, cell.getContext())} + + ))} + + )) + ) : ( + + +
+

No keys found

+
+ + + )} + +
+
+
+
+
+ )} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/team/tabVisibilityUtils.test.ts b/ui/litellm-dashboard/src/components/team/tabVisibilityUtils.test.ts index 5a09b4fa36c..8b9a0402c9a 100644 --- a/ui/litellm-dashboard/src/components/team/tabVisibilityUtils.test.ts +++ b/ui/litellm-dashboard/src/components/team/tabVisibilityUtils.test.ts @@ -11,6 +11,7 @@ describe("team_info_tabs", () => { describe("TEAM_INFO_TAB_LABELS", () => { it("should have label for every tab key", () => { expect(TEAM_INFO_TAB_LABELS[TEAM_INFO_TAB_KEYS.OVERVIEW]).toBe("Overview"); + expect(TEAM_INFO_TAB_LABELS[TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS]).toBe("Virtual Keys"); expect(TEAM_INFO_TAB_LABELS[TEAM_INFO_TAB_KEYS.MEMBERS]).toBe("Members"); expect(TEAM_INFO_TAB_LABELS[TEAM_INFO_TAB_KEYS.MEMBER_PERMISSIONS]).toBe("Member Permissions"); expect(TEAM_INFO_TAB_LABELS[TEAM_INFO_TAB_KEYS.SETTINGS]).toBe("Settings"); @@ -18,15 +19,16 @@ describe("team_info_tabs", () => { }); describe("getTeamInfoVisibleTabs", () => { - it("returns only overview when user cannot edit team", () => { + it("returns overview and virtual keys when user cannot edit team", () => { const tabs = getTeamInfoVisibleTabs(false); - expect(tabs).toEqual([TEAM_INFO_TAB_KEYS.OVERVIEW]); + expect(tabs).toEqual([TEAM_INFO_TAB_KEYS.OVERVIEW, TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS]); }); it("returns all tabs when user can edit team", () => { const tabs = getTeamInfoVisibleTabs(true); expect(tabs).toEqual([ TEAM_INFO_TAB_KEYS.OVERVIEW, + TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS, TEAM_INFO_TAB_KEYS.MEMBERS, TEAM_INFO_TAB_KEYS.MEMBER_PERMISSIONS, TEAM_INFO_TAB_KEYS.SETTINGS, @@ -55,6 +57,19 @@ describe("team_info_tabs", () => { expect(isTeamInfoTabVisible(TEAM_INFO_TAB_KEYS.OVERVIEW, true)).toBe(true); }); + it("always returns true for virtual keys tab regardless of edit permission", () => { + expect(isTeamInfoTabVisible(TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS, false)).toBe(true); + expect(isTeamInfoTabVisible(TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS, true)).toBe(true); + }); + + it("returns false for member permissions tab when user cannot edit", () => { + expect(isTeamInfoTabVisible(TEAM_INFO_TAB_KEYS.MEMBER_PERMISSIONS, false)).toBe(false); + }); + + it("returns true for member permissions tab when user can edit", () => { + expect(isTeamInfoTabVisible(TEAM_INFO_TAB_KEYS.MEMBER_PERMISSIONS, true)).toBe(true); + }); + it("returns false for members tab when user cannot edit", () => { expect(isTeamInfoTabVisible(TEAM_INFO_TAB_KEYS.MEMBERS, false)).toBe(false); }); diff --git a/ui/litellm-dashboard/src/components/team/tabVisibilityUtils.ts b/ui/litellm-dashboard/src/components/team/tabVisibilityUtils.ts index d77230ea09b..dd0e54baf36 100644 --- a/ui/litellm-dashboard/src/components/team/tabVisibilityUtils.ts +++ b/ui/litellm-dashboard/src/components/team/tabVisibilityUtils.ts @@ -5,6 +5,7 @@ export const TEAM_INFO_TAB_KEYS = { OVERVIEW: "overview", + VIRTUAL_KEYS: "virtual-keys", MEMBERS: "members", MEMBER_PERMISSIONS: "member-permissions", SETTINGS: "settings", @@ -12,6 +13,7 @@ export const TEAM_INFO_TAB_KEYS = { export const TEAM_INFO_TAB_LABELS: Record = { [TEAM_INFO_TAB_KEYS.OVERVIEW]: "Overview", + [TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS]: "Virtual Keys", [TEAM_INFO_TAB_KEYS.MEMBERS]: "Members", [TEAM_INFO_TAB_KEYS.MEMBER_PERMISSIONS]: "Member Permissions", [TEAM_INFO_TAB_KEYS.SETTINGS]: "Settings", @@ -19,11 +21,11 @@ export const TEAM_INFO_TAB_LABELS: Record = { /** * Returns the list of tab keys that should be visible based on permissions. - * - Overview: always visible + * - Overview, Virtual Keys: always visible * - Members, Member Permissions, Settings: only when canEditTeam is true */ export function getTeamInfoVisibleTabs(canEditTeam: boolean): readonly string[] { - const baseTabs = [TEAM_INFO_TAB_KEYS.OVERVIEW]; + const baseTabs = [TEAM_INFO_TAB_KEYS.OVERVIEW, TEAM_INFO_TAB_KEYS.VIRTUAL_KEYS]; if (canEditTeam) { return [ ...baseTabs, From 6d98622923f95c6e047c910637cdf92b3f055f77 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Mon, 23 Feb 2026 15:21:49 -0800 Subject: [PATCH 005/147] filters working properly for team virtual keys --- .../key_team_helpers/filter_helpers.ts | 77 +++++++- .../src/components/team/TeamInfo.test.tsx | 71 +++++++- .../team/TeamVirtualKeysTable.test.tsx | 170 +++++++++++++++++- .../components/team/TeamVirtualKeysTable.tsx | 133 ++++++++------ 4 files changed, 390 insertions(+), 61 deletions(-) diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts index fa12a75aacd..08eccc4ed2c 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts @@ -1,7 +1,82 @@ -import { teamListCall, organizationListCall, keyAliasesCall } from "../networking" +import { teamListCall, organizationListCall, keyAliasesCall, keyListCall } from "../networking"; import { Team } from "./key_list"; import { Organization } from "../networking"; +export interface TeamFilterOptions { + keyAliases: string[]; + organizationIds: string[]; + userIds: Array<{ id: string; email: string }>; +} + +/** + * Fetches filter options (key aliases, org IDs, user IDs) scoped to a team's keys. + * Used by TeamVirtualKeysTable to show only relevant filter options. + */ +export const fetchTeamFilterOptions = async ( + accessToken: string | null, + teamId: string, +): Promise => { + if (!accessToken || !teamId) { + return { keyAliases: [], organizationIds: [], userIds: [] }; + } + + try { + const keyAliases = new Set(); + const organizationIds = new Set(); + const userMap = new Map(); // user_id -> user_email + + let page = 1; + let totalPages = 1; + + do { + const response = await keyListCall( + accessToken, + null, + teamId, + null, + null, + null, + page, + 100, + null, + null, + "user", + null, + ); + + const keys = response?.keys || []; + totalPages = response?.total_pages ?? 1; + + for (const key of keys) { + const alias = key?.key_alias; + if (alias && typeof alias === "string") { + keyAliases.add(alias.trim()); + } + const orgId = key?.organization_id; + if (orgId && typeof orgId === "string") { + organizationIds.add(orgId.trim()); + } + const userId = key?.user_id; + if (userId && typeof userId === "string") { + const email = key?.user?.user_email || userId; + userMap.set(userId, email); + } + } + + page++; + } while (page <= totalPages); + + return { + keyAliases: Array.from(keyAliases).sort(), + organizationIds: Array.from(organizationIds).sort(), + userIds: Array.from(userMap.entries()).map(([id, email]) => ({ id, email })), + }; + } catch (error) { + console.error("Error fetching team filter options:", error); + return { keyAliases: [], organizationIds: [], userIds: [] }; + } +}; + /** * Fetches all key aliases via the dedicated /key/aliases endpoint * @param accessToken The access token for API authentication diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 75fe76902bd..317db7b37fd 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -110,6 +110,11 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ })); vi.mock("../key_team_helpers/filter_helpers", () => ({ + fetchTeamFilterOptions: vi.fn().mockResolvedValue({ + keyAliases: [], + organizationIds: [], + userIds: [], + }), fetchAllKeyAliases: vi.fn().mockResolvedValue([]), fetchAllOrganizations: vi.fn().mockResolvedValue([]), })); @@ -597,8 +602,22 @@ describe("TeamInfoView", () => { it("should display X Members in Virtual Keys tab when navigated to", async () => { const user = userEvent.setup(); vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + const fiveKeys = Array.from({ length: 5 }, (_, i) => ({ + token: `sk-${i}`, + token_id: `key-${i}`, + key_alias: `key_${i}`, + key_name: `sk-...${i}`, + user_id: `user-${i}`, + organization_id: null, + user: { user_id: `user-${i}`, user_email: `user${i}@test.com` }, + created_at: "2024-01-01T00:00:00Z", + team_id: "123", + spend: 0, + max_budget: 100, + models: ["gpt-4"], + })); mockUseKeys.mockReturnValue({ - data: { keys: [], total_count: 5, current_page: 1, total_pages: 1 }, + data: { keys: fiveKeys, total_count: 5, current_page: 1, total_pages: 1 }, isPending: false, isFetching: false, refetch: vi.fn(), @@ -619,6 +638,56 @@ describe("TeamInfoView", () => { }); }); + it("should show Filters and pagination controls in Virtual Keys tab", async () => { + const user = userEvent.setup(); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + mockUseKeys.mockReturnValue({ + data: { + keys: [ + { + token: "sk-1", + token_id: "key-1", + key_alias: "key1", + key_name: "sk-...1", + user_id: "user-1", + organization_id: null, + user: { user_id: "user-1", user_email: "user1@test.com" }, + created_at: "2024-01-01T00:00:00Z", + team_id: "123", + spend: 0, + max_budget: 100, + models: ["gpt-4"], + }, + ], + total_count: 1, + current_page: 1, + total_pages: 1, + }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const virtualKeysTab = screen.getByRole("tab", { name: "Virtual Keys" }); + await user.click(virtualKeysTab); + + await waitFor(() => { + expect(screen.getByText("1 Member")).toBeInTheDocument(); + }); + expect(screen.getByRole("button", { name: "Filters" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Reset Filters" })).toBeInTheDocument(); + expect(screen.getByText("Page 1 of 1")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Previous" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Next" })).toBeInTheDocument(); + }); + it("should display object permissions when present", async () => { vi.mocked(networking.teamInfoCall).mockResolvedValue( createMockTeamData({ diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx index e4b7b00fb58..41df5611e07 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -4,6 +4,7 @@ import { beforeEach, describe, expect, it, vi, MockedFunction } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { fetchTeamFilterOptions } from "../key_team_helpers/filter_helpers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { KeyResponse } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; @@ -17,8 +18,11 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ })); vi.mock("../key_team_helpers/filter_helpers", () => ({ - fetchAllKeyAliases: vi.fn().mockResolvedValue([]), - fetchAllOrganizations: vi.fn().mockResolvedValue([]), + fetchTeamFilterOptions: vi.fn().mockResolvedValue({ + keyAliases: [], + organizationIds: [], + userIds: [], + }), })); vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({ @@ -139,7 +143,7 @@ describe("TeamVirtualKeysTable", () => { }); }); - it("should call useKeys with expand user to fetch user email", async () => { + it("should call useKeys with page, pageSize, and expand user for server-side pagination", async () => { renderWithProviders(); await waitFor(() => { @@ -189,4 +193,164 @@ describe("TeamVirtualKeysTable", () => { }); expect(screen.getByText("Key ID")).toBeInTheDocument(); }); + + it("should display keys in table when data is loaded", async () => { + mockUseKeys.mockReturnValue({ + data: { + keys: [ + createMockKey({ key_alias: "alice_key_team1" }), + createMockKey({ token: "sk-2", token_id: "key-2", key_alias: "bob_key_team1" }), + ], + total_count: 2, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("2 Members")).toBeInTheDocument(); + }); + expect(screen.getByText("alice_key_team1")).toBeInTheDocument(); + expect(screen.getByText("bob_key_team1")).toBeInTheDocument(); + }); + + it("should show Page X of Y when multiple pages exist", async () => { + mockUseKeys.mockReturnValue({ + data: { + keys: [createMockKey()], + total_count: 100, + current_page: 1, + total_pages: 3, + } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Page 1 of 3")).toBeInTheDocument(); + }); + expect(screen.getByText("100 Members")).toBeInTheDocument(); + }); + + it("should fetch page 2 when Next is clicked", async () => { + const user = userEvent.setup(); + mockUseKeys.mockImplementation((page: number) => ({ + data: { + keys: page === 1 ? [createMockKey()] : [createMockKey({ token: "sk-page2", key_alias: "page2_key" })], + total_count: 100, + current_page: page, + total_pages: 3, + } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any)); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Page 1 of 3")).toBeInTheDocument(); + }); + + const nextButton = screen.getByRole("button", { name: "Next" }); + await user.click(nextButton); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 2, + 50, + expect.objectContaining({ teamID: "team-1" }) + ); + }); + }); + + it("should show Loading keys when isPending", async () => { + mockUseKeys.mockReturnValue({ + data: undefined, + isPending: true, + isFetching: true, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Loading keys...")).toBeInTheDocument(); + }); + }); + + it("should show No keys found when keys array is empty", async () => { + mockUseKeys.mockReturnValue({ + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("0 Members")).toBeInTheDocument(); + }); + expect(screen.getByText("No keys found")).toBeInTheDocument(); + }); + + it("should fetch team-scoped filter options for Key Alias, Organization ID, and User ID", async () => { + const mockFetchTeamFilterOptions = vi.mocked(fetchTeamFilterOptions); + mockFetchTeamFilterOptions.mockResolvedValue({ + keyAliases: ["alice_key_team1", "charlie_key_team1"], + organizationIds: ["org-123"], + userIds: [ + { id: "user-1", email: "alice@example.com" }, + { id: "user-2", email: "charlie@example.com" }, + ], + }); + + // Use unique teamId to avoid cache hit from previous tests (refetchOnMount: false) + renderWithProviders( + + ); + + await waitFor(() => { + expect(mockFetchTeamFilterOptions).toHaveBeenCalledWith( + "test-token", + "team-filter-options-test" + ); + }); + }); + + it("should open Key Info View when key is clicked", async () => { + mockUseKeys.mockReturnValue({ + data: { + keys: [createMockKey({ token: "sk-click-me", key_alias: "clickable_key" })], + total_count: 1, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("clickable_key")).toBeInTheDocument(); + }); + + const keyButton = screen.getByRole("button", { name: /sk-click-me|clickable_key/ }); + await userEvent.click(keyButton); + + await waitFor(() => { + expect(screen.getByText("Key Info View")).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index ce998d18677..ddbffa118b1 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -8,7 +8,6 @@ import { ColumnDef, flexRender, getCoreRowModel, - getPaginationRowModel, getSortedRowModel, PaginationState, SortingState, @@ -28,14 +27,14 @@ import { } from "@tremor/react"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Popover, Skeleton, Tooltip } from "antd"; -import React, { useEffect, useMemo, useState } from "react"; +import React, { useCallback, useEffect, useMemo, useState } from "react"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import FilterComponent, { FilterOption } from "../molecules/filter"; import { Organization } from "../networking"; import KeyInfoView from "../templates/key_info_view"; import { useQuery } from "@tanstack/react-query"; -import { fetchAllKeyAliases, fetchAllOrganizations } from "../key_team_helpers/filter_helpers"; +import { fetchTeamFilterOptions } from "../key_team_helpers/filter_helpers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; interface TeamVirtualKeysTableProps { @@ -69,22 +68,24 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi const sortBy = sorting.length > 0 ? sorting[0].id : "created_at"; const sortOrder = sorting.length > 0 ? (sorting[0].desc ? "desc" : "asc") : "desc"; + const pageIndex = tablePagination.pageIndex; + const pageSize = tablePagination.pageSize; + const { data: keys, isPending: isLoading, isFetching, refetch, - } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, { + } = useKeys(pageIndex + 1, pageSize, { teamID: teamId, + organizationID: filters["Organization ID"]?.trim() || undefined, + selectedKeyAlias: filters["Key Alias"]?.trim() || undefined, + userID: filters["User ID"]?.trim() || undefined, sortBy: sortBy || undefined, sortOrder: sortOrder || undefined, - organizationID: filters["Organization ID"] || undefined, - selectedKeyAlias: filters["Key Alias"] || undefined, - userID: filters["User ID"] || undefined, expand: "user", }); - const totalCount = keys?.total_count || 0; const displayKeys = useMemo(() => { const kList = keys?.keys || []; const orgId = organization?.organization_id; @@ -94,6 +95,9 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi organization_id: k.organization_id || orgId, })); }, [keys?.keys, organization?.organization_id]); + + const totalCount = keys?.total_count ?? 0; + const pageCount = keys?.total_pages ?? 0; const [expandedAccordions, setExpandedAccordions] = useState>({}); const currentTeam: Team = useMemo( @@ -114,29 +118,27 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi [teamId, teamAlias, organization], ); - const allKeyAliasesQuery = useQuery({ - queryKey: ["allKeyAliases"], - queryFn: async () => fetchAllKeyAliases(accessToken), - enabled: !!accessToken, + const teamFilterOptionsQuery = useQuery({ + queryKey: ["teamFilterOptions", teamId], + queryFn: async () => fetchTeamFilterOptions(accessToken, teamId), + enabled: !!accessToken && !!teamId, }); - const allKeyAliases = allKeyAliasesQuery.data || []; + const teamFilterOptions = teamFilterOptionsQuery.data || { + keyAliases: [], + organizationIds: [], + userIds: [], + }; - const allOrganizationsQuery = useQuery({ - queryKey: ["allOrganizations"], - queryFn: async () => fetchAllOrganizations(accessToken), - enabled: !!accessToken, - }); - const allOrganizations = allOrganizationsQuery.data || []; - - useEffect(() => { - if (refetch) { - const handleStorageChange = () => refetch(); - window.addEventListener("storage", handleStorageChange); - return () => window.removeEventListener("storage", handleStorageChange); - } + const handleStorageChange = useCallback(() => { + refetch?.(); }, [refetch]); - const handleFilterChange = (newFilters: Record, skipDebounce = false) => { + useEffect(() => { + window.addEventListener("storage", handleStorageChange); + return () => window.removeEventListener("storage", handleStorageChange); + }, [handleStorageChange]); + + const handleFilterChange = useCallback((newFilters: Record, skipDebounce = false) => { setFilters((prev) => ({ ...prev, "Organization ID": newFilters["Organization ID"] ?? prev["Organization ID"], @@ -148,9 +150,9 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi if (!skipDebounce) { setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); } - }; + }, []); - const handleFilterReset = () => { + const handleFilterReset = useCallback(() => { setFilters({ "Organization ID": "", "Key Alias": "", @@ -159,7 +161,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi "Sort Order": "desc", }); setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }; + }, []); const filterOptions: FilterOption[] = useMemo( () => [ @@ -168,16 +170,13 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi label: "Organization ID", isSearchable: true, searchFn: async (searchText: string) => { - if (!allOrganizations.length) return []; - const filtered = allOrganizations.filter( - (org) => org.organization_id?.toLowerCase().includes(searchText.toLowerCase()) ?? false, - ); - return filtered - .filter((org) => org.organization_id != null) - .map((org) => ({ - label: `${org.organization_id || "Unknown"} (${org.organization_id})`, - value: org.organization_id as string, - })); + const { organizationIds } = teamFilterOptions; + if (!organizationIds.length) return []; + const lower = searchText.toLowerCase(); + const filtered = lower + ? organizationIds.filter((id) => id.toLowerCase().includes(lower)) + : organizationIds; + return filtered.map((id) => ({ label: id, value: id })); }, }, { @@ -185,15 +184,35 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi label: "Key Alias", isSearchable: true, searchFn: async (searchText: string) => { - const filtered = allKeyAliases.filter((alias) => - alias.toLowerCase().includes(searchText.toLowerCase()), - ); + const { keyAliases } = teamFilterOptions; + const lower = searchText.toLowerCase(); + const filtered = lower + ? keyAliases.filter((alias) => alias.toLowerCase().includes(lower)) + : keyAliases; return filtered.map((alias) => ({ label: alias, value: alias })); }, }, - { name: "User ID", label: "User ID", isSearchable: false }, + { + name: "User ID", + label: "User ID", + isSearchable: true, + searchFn: async (searchText: string) => { + const { userIds } = teamFilterOptions; + const lower = searchText.toLowerCase(); + const filtered = lower + ? userIds.filter( + (u) => + u.id.toLowerCase().includes(lower) || u.email.toLowerCase().includes(lower), + ) + : userIds; + return filtered.map((u) => ({ + label: u.email ? `${u.id} (${u.email})` : u.id, + value: u.id, + })); + }, + }, ], - [allOrganizations, allKeyAliases], + [teamFilterOptions], ); const columns: ColumnDef[] = useMemo( @@ -521,13 +540,8 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi [expandedAccordions], ); - const table = useReactTable({ - data: displayKeys, - columns, - columnResizeMode: "onChange", - columnResizeDirection: "ltr", - state: { sorting, pagination: tablePagination }, - onSortingChange: (updaterOrValue) => { + const handleSortingChange = useCallback( + (updaterOrValue: React.SetStateAction) => { const newSorting = typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue; setSorting(newSorting); @@ -535,7 +549,6 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi const sortState = newSorting[0]; handleFilterChange( { - ...filters, "Sort By": sortState.id, "Sort Order": sortState.desc ? "desc" : "asc", }, @@ -543,17 +556,25 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi ); } }, + [sorting, handleFilterChange], + ); + + const table = useReactTable({ + data: displayKeys, + columns, + columnResizeMode: "onChange", + columnResizeDirection: "ltr", + state: { sorting, pagination: tablePagination }, + onSortingChange: handleSortingChange, onPaginationChange: setTablePagination, getCoreRowModel: getCoreRowModel(), getSortedRowModel: getSortedRowModel(), - getPaginationRowModel: getPaginationRowModel(), enableSorting: true, manualSorting: false, manualPagination: true, - pageCount: Math.ceil(totalCount / tablePagination.pageSize), + pageCount: pageCount, }); - const { pageIndex, pageSize } = table.getState().pagination; return (
{selectedKey ? ( From f44c36f9802bc46bbd4b646924aed8ddeea82f2a Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Mon, 23 Feb 2026 15:43:04 -0800 Subject: [PATCH 006/147] Fix Unbounded pagination --- .../src/components/key_team_helpers/filter_helpers.ts | 6 ++++-- .../src/components/team/TeamVirtualKeysTable.tsx | 1 + 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts index 08eccc4ed2c..dd3766592bf 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts @@ -25,6 +25,8 @@ export const fetchTeamFilterOptions = async ( const organizationIds = new Set(); const userMap = new Map(); // user_id -> user_email + const MAX_PAGES = 1; // Cap at 20 keys on load to avoid heavy fetches + const PAGE_SIZE = 20; let page = 1; let totalPages = 1; @@ -37,7 +39,7 @@ export const fetchTeamFilterOptions = async ( null, null, page, - 100, + PAGE_SIZE, null, null, "user", @@ -64,7 +66,7 @@ export const fetchTeamFilterOptions = async ( } page++; - } while (page <= totalPages); + } while (page <= totalPages && page <= MAX_PAGES); return { keyAliases: Array.from(keyAliases).sort(), diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index ddbffa118b1..13d7885d672 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -122,6 +122,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi queryKey: ["teamFilterOptions", teamId], queryFn: async () => fetchTeamFilterOptions(accessToken, teamId), enabled: !!accessToken && !!teamId, + staleTime: 30000, // 30 seconds - align with useKeys }); const teamFilterOptions = teamFilterOptionsQuery.data || { keyAliases: [], From bc99e38e55348aaab5db5fcef3adb1f9ba44faa3 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Mon, 23 Feb 2026 15:44:51 -0800 Subject: [PATCH 007/147] Fix: accessToken not included in queryKey --- .../src/components/team/TeamVirtualKeysTable.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 13d7885d672..74100069779 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -119,7 +119,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi ); const teamFilterOptionsQuery = useQuery({ - queryKey: ["teamFilterOptions", teamId], + queryKey: ["teamFilterOptions", teamId, accessToken], queryFn: async () => fetchTeamFilterOptions(accessToken, teamId), enabled: !!accessToken && !!teamId, staleTime: 30000, // 30 seconds - align with useKeys From a1f5450b4001555b04ce51044004580ced8e3cb1 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Mon, 23 Feb 2026 16:01:04 -0800 Subject: [PATCH 008/147] manualSorting fix. Filtering fetching fix --- .../components/key_team_helpers/filter_helpers.ts | 14 ++++++++------ .../src/components/team/TeamVirtualKeysTable.tsx | 2 +- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts index dd3766592bf..0a5590d78ce 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts @@ -8,9 +8,13 @@ export interface TeamFilterOptions { userIds: Array<{ id: string; email: string }>; } +const FILTER_OPTIONS_PAGE_SIZE = 100; // API max per page +const MAX_PAGES = 50; // Cap at 5000 keys to avoid unbounded fetches + /** - * Fetches filter options (key aliases, org IDs, user IDs) scoped to a team's keys. - * Used by TeamVirtualKeysTable to show only relevant filter options. + * Fetches filter options (key aliases, org IDs, user IDs) from all team keys. + * Paginates through pages to build complete dropdowns. Capped at 50 pages + * (5000 keys) to limit load for very large teams. */ export const fetchTeamFilterOptions = async ( accessToken: string | null, @@ -23,10 +27,8 @@ export const fetchTeamFilterOptions = async ( try { const keyAliases = new Set(); const organizationIds = new Set(); - const userMap = new Map(); // user_id -> user_email + const userMap = new Map(); - const MAX_PAGES = 1; // Cap at 20 keys on load to avoid heavy fetches - const PAGE_SIZE = 20; let page = 1; let totalPages = 1; @@ -39,7 +41,7 @@ export const fetchTeamFilterOptions = async ( null, null, page, - PAGE_SIZE, + FILTER_OPTIONS_PAGE_SIZE, null, null, "user", diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 74100069779..de3061a003d 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -571,7 +571,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi getCoreRowModel: getCoreRowModel(), getSortedRowModel: getSortedRowModel(), enableSorting: true, - manualSorting: false, + manualSorting: true, // Server sorts via useKeys. Avoid redundant client-side sort manualPagination: true, pageCount: pageCount, }); From 18bc0d4ec94f04bc3b284003f43a90f6d15a39e4 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Fri, 27 Feb 2026 13:37:23 -0800 Subject: [PATCH 009/147] reusable-credentials --- .../handle_add_model_submit.test.tsx | 24 +++++++ .../add_model/handle_add_model_submit.tsx | 3 + .../src/components/model_info_view.test.tsx | 63 +++++++++++++++++++ .../src/components/model_info_view.tsx | 57 ++++++++++++++++- 4 files changed, 146 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx index 5cf7a18634e..c261e505684 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx @@ -55,4 +55,28 @@ describe("prepareModelAddRequest", () => { const [deployment] = deployments!; expect(deployment.litellmParamsObj.custom_llm_provider).toBe("petals"); }); + + it("ignores litellm_credential_name inside LiteLLM Params JSON", async () => { + const formValues = { + model_mappings: [ + { + public_name: "Public Model", + litellm_model: "litellm/public", + }, + ], + model_name: "custom-model-name", + litellm_credential_name: "selected-credential", + litellm_extra_params: JSON.stringify({ + litellm_credential_name: "from-json", + timeout: 5, + }), + }; + + const deployments = await prepareModelAddRequest({ ...formValues }, "token", null); + + expect(deployments).toHaveLength(1); + const [deployment] = deployments!; + expect(deployment.litellmParamsObj.litellm_credential_name).toBe("selected-credential"); + expect(deployment.litellmParamsObj.timeout).toBe(5); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx index 1d8c980c5ae..5137a302dd5 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx @@ -91,6 +91,9 @@ export const prepareModelAddRequest = async (formValues: Record, ac if (value && value != undefined) { try { litellmExtraParams = JSON.parse(value); + if ("litellm_credential_name" in litellmExtraParams) { + delete litellmExtraParams.litellm_credential_name; + } } catch (error) { NotificationManager.fromBackend("Failed to parse LiteLLM Extra Params: " + error); throw new Error("Failed to parse litellm_extra_params: " + error); diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 7158c452d94..b2cbbda34aa 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -23,6 +23,7 @@ vi.mock("./molecules/notifications_manager", () => ({ vi.mock("./networking", () => ({ modelInfoV1Call: vi.fn(), credentialGetCall: vi.fn(), + credentialListCall: vi.fn(), getGuardrailsList: vi.fn(), tagListCall: vi.fn(), testConnectionRequest: vi.fn(), @@ -47,6 +48,7 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ const mockNotificationsManager = vi.mocked(NotificationsManager); const mockModelInfoV1Call = vi.mocked(networking.modelInfoV1Call); const mockCredentialGetCall = vi.mocked(networking.credentialGetCall); +const mockCredentialListCall = vi.mocked(networking.credentialListCall); const mockGetGuardrailsList = vi.mocked(networking.getGuardrailsList); const mockTagListCall = vi.mocked(networking.tagListCall); const mockTestConnectionRequest = vi.mocked(networking.testConnectionRequest); @@ -63,6 +65,7 @@ describe("ModelInfoView", () => { model: "gpt-4", api_base: "https://api.openai.com/v1", custom_llm_provider: "openai", + litellm_credential_name: "selected-credential", }, model_info: { id: "123", @@ -125,6 +128,15 @@ describe("ModelInfoView", () => { credential_values: {}, credential_info: {}, }); + mockCredentialListCall.mockResolvedValue({ + credentials: [ + { + credential_name: "selected-credential", + credential_values: {}, + credential_info: {}, + }, + ], + }); mockGetGuardrailsList.mockResolvedValue({ guardrails: [{ guardrail_name: "content_filter" }, { guardrail_name: "toxicity_filter" }], @@ -489,6 +501,57 @@ describe("ModelInfoView", () => { }); }); + it("should show existing credentials field in edit mode", async () => { + const user = userEvent.setup(); + render(, { wrapper }); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await waitFor(() => { + expect(screen.getByText("Existing Credentials")).toBeInTheDocument(); + }); + }); + + it("should keep selector credential and ignore litellm_credential_name from LiteLLM Params json", async () => { + const user = userEvent.setup(); + render(, { wrapper }); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + const litellmParamsInput = screen + .getAllByRole("textbox") + .find( + (input) => + input.tagName === "TEXTAREA" && + (input as HTMLTextAreaElement).value.includes('"custom_llm_provider"'), + ); + expect(litellmParamsInput).toBeDefined(); + if (!litellmParamsInput) { + return; + } + expect((litellmParamsInput as HTMLTextAreaElement).value).not.toContain("litellm_credential_name"); + await user.clear(litellmParamsInput); + await user.paste(`{"litellm_credential_name":"from-json","timeout":42}`); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(mockModelPatchUpdateCall).toHaveBeenCalled(); + }); + + const updatePayload = mockModelPatchUpdateCall.mock.calls[0][1]; + expect(updatePayload.litellm_params.litellm_credential_name).toBe("selected-credential"); + expect(updatePayload.litellm_params.litellm_credential_name).not.toBe("from-json"); + }); + it("should display health check model field for wildcard models", async () => { const wildcardModelData = { ...defaultModelData, diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index e2fc8caa21c..0ba86e9f243 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -30,6 +30,7 @@ import { CredentialItem, credentialCreateCall, credentialGetCall, + credentialListCall, getGuardrailsList, modelDeleteCall, modelInfoV1Call, @@ -75,6 +76,7 @@ export default function ModelInfoView({ const [isAutoRouterModalOpen, setIsAutoRouterModalOpen] = useState(false); const [guardrailsList, setGuardrailsList] = useState([]); const [tagsList, setTagsList] = useState>({}); + const [credentialsList, setCredentialsList] = useState([]); // Fetch model data using hook const { data: rawModelDataResponse, isLoading: isLoadingModel } = useModelsInfo(1, 50, undefined, modelId); @@ -191,10 +193,21 @@ export default function ModelInfoView({ } }; + const fetchCredentials = async () => { + if (!accessToken) return; + try { + const response = await credentialListCall(accessToken); + setCredentialsList(response.credentials || []); + } catch (error) { + console.error("Failed to fetch credentials:", error); + } + }; + getExistingCredential(); getModelInfo(); fetchGuardrails(); fetchTags(); + fetchCredentials(); }, [accessToken, modelId]); const handleReuseCredential = async (values: any) => { @@ -220,6 +233,7 @@ export default function ModelInfoView({ let parsedExtraParams: Record = {}; try { parsedExtraParams = values.litellm_extra_params ? JSON.parse(values.litellm_extra_params) : {}; + delete parsedExtraParams.litellm_credential_name; } catch (e) { NotificationsManager.fromBackend("Invalid JSON in LiteLLM Params"); setIsSaving(false); @@ -242,6 +256,11 @@ export default function ModelInfoView({ output_cost_per_token: values.output_cost / 1_000_000, tags: values.tags, }; + if (values.litellm_credential_name) { + updatedLitellmParams.litellm_credential_name = values.litellm_credential_name; + } else { + delete updatedLitellmParams.litellm_credential_name; + } if (values.guardrails) { updatedLitellmParams.guardrails = values.guardrails; } @@ -608,7 +627,16 @@ export default function ModelInfoView({ : [], tags: Array.isArray(localModelData.litellm_params?.tags) ? localModelData.litellm_params.tags : [], health_check_model: isWildcardModel ? localModelData.model_info?.health_check_model : null, - litellm_extra_params: JSON.stringify(localModelData.litellm_params || {}, null, 2), + litellm_credential_name: localModelData.litellm_params?.litellm_credential_name || undefined, + litellm_extra_params: JSON.stringify( + Object.fromEntries( + Object.entries(localModelData.litellm_params || {}).filter( + ([key]) => key !== "litellm_credential_name", + ), + ), + null, + 2, + ), }} layout="vertical" onValuesChange={() => setIsDirty(true)} @@ -930,6 +958,33 @@ export default function ModelInfoView({
)} +
+ Existing Credentials + {isEditing ? ( + + setObjectIdSearch(e.target.value)} - className="px-3 py-2 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500" - /> -
- - - - - -
- {/* Custom Action Filter Dropdown */} -
- - - {actionFilterOpen && ( -
-
- {[ - { label: "All Actions", value: "all" }, - { label: "Created", value: "created" }, - { label: "Updated", value: "updated" }, - { label: "Deleted", value: "deleted" }, - { label: "Rotated", value: "rotated" }, - ].map((option) => ( - - ))} -
-
- )} -
- - {/* Custom Table Filter Dropdown */} -
- - - {tableFilterOpen && ( -
-
- {[ - { label: "All Tables", value: "all" }, - { label: "Keys", value: "keys" }, - { label: "Teams", value: "teams" }, - { label: "Users", value: "users" }, - ].map((option) => ( - - ))} -
-
- )} -
- - - Showing {allLogsQuery.isLoading ? "..." : currentDisplayItemsStart} -{" "} - {allLogsQuery.isLoading ? "..." : currentDisplayItemsEnd} of{" "} - {allLogsQuery.isLoading ? "..." : totalFilteredItems} results - -
- - Page {allLogsQuery.isLoading ? "..." : clientCurrentPage} of{" "} - {allLogsQuery.isLoading ? "..." : totalFilteredPages} - - - -
-
+ {/* Filters */} +
+ { + setObjectId(val); + handleFilterChange(); + }} + onChange={(e) => { + if (!e.target.value) { + setObjectId(""); + handleFilterChange(); + } + }} + /> + { + setChangedBy(val); + handleFilterChange(); + }} + onChange={(e) => { + if (!e.target.value) { + setChangedBy(""); + handleFilterChange(); + } + }} + /> + { + setTeamId(val); + handleFilterChange(); + }} + onChange={(e) => { + if (!e.target.value) { + setTeamId(""); + handleFilterChange(); + } + }} + /> + { + setKeyHash(val); + handleFilterChange(); + }} + onChange={(e) => { + if (!e.target.value) { + setKeyHash(""); + handleFilterChange(); + } + }} + /> + { + setTableName(val); + handleFilterChange(); + }} + />
- true} + + {/* Table */} + + columns={columns} + dataSource={auditLogs} + rowKey="id" + loading={query.isLoading} + size="small" + onRow={(record) => ({ + onClick: () => handleRowClick(record), + style: { cursor: "pointer" }, + })} + pagination={{ + current: page, + pageSize: PAGE_SIZE, + total, + showTotal: (t) => `${t} total`, + showSizeChanger: false, + onChange: (p) => setPage(p), + }} + onChange={handleTableChange} /> + + setDrawerOpen(false)} + log={selectedLog} + /> ); } From 5655cb87fc7e95b68e70f4018fab5ce4b0b5d0a8 Mon Sep 17 00:00:00 2001 From: openhands Date: Mon, 2 Mar 2026 08:11:07 +0000 Subject: [PATCH 015/147] fix: pass all custom pricing fields to register_model in completion() and embedding() Previously, register_model() was called with only input_cost_per_token, output_cost_per_token, and litellm_provider. This dropped ~40+ other pricing fields from CustomPricingLiteLLMParams (cache_read_input_token_cost, cache_creation_input_token_cost, output_cost_per_reasoning_token, etc.) as well as model_info metadata (mode, supports_prompt_caching, max_tokens). For DB-sourced custom-priced models, the first request after a pod restart would register a partial entry in litellm.model_cost, causing cost calculations to miss cache token discounts and other extended pricing until the entry was later enriched by deployment_callback_on_success mutating the lru_cache. Changes: - Add _build_custom_pricing_entry() helper that iterates over all CustomPricingLiteLLMParams.model_fields and merges model_info metadata - Replace hardcoded 3-field dicts in both completion() and embedding() with the new helper - Add 7 tests covering field collection, model_info merging, precedence, None skipping, and end-to-end register_model behavior Co-authored-by: openhands --- litellm/main.py | 80 ++++---- .../test_register_model_custom_pricing.py | 180 ++++++++++++++++++ 2 files changed, 223 insertions(+), 37 deletions(-) create mode 100644 tests/test_litellm/test_register_model_custom_pricing.py diff --git a/litellm/main.py b/litellm/main.py index cb3ddc2f401..737cf50ad24 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -107,6 +107,7 @@ from litellm.realtime_api.main import _realtime_health_check from litellm.secret_managers.main import get_secret_bool, get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + CustomPricingLiteLLMParams, ModelResponseStream, RawRequestTypedDict, StreamingChoices, @@ -996,6 +997,32 @@ def _drop_input_examples_from_tools( return cleaned_tools +def _build_custom_pricing_entry( + custom_llm_provider: str, + kwargs: dict, + model_info: Optional[dict] = None, +) -> dict: + """Build a complete model cost entry from kwargs and model_info. + + Collects all CustomPricingLiteLLMParams fields present in kwargs and + merges metadata from model_info (mode, supports_prompt_caching, max_tokens) + so that register_model() receives the full pricing configuration. + """ + entry: dict = {"litellm_provider": custom_llm_provider} + + for field_name in CustomPricingLiteLLMParams.model_fields: + value = kwargs.get(field_name) + if value is not None: + entry[field_name] = value + + if model_info and isinstance(model_info, dict): + for key in ("mode", "supports_prompt_caching", "max_tokens"): + if key in model_info and model_info[key] is not None: + entry.setdefault(key, model_info[key]) + + return entry + + @tracer.wrap() @client def completion( # type: ignore # noqa: PLR0915 @@ -1351,27 +1378,16 @@ def completion( # type: ignore # noqa: PLR0915 timeout = float(timeout) # type: ignore ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### - if input_cost_per_token is not None and output_cost_per_token is not None: + if ( + input_cost_per_token is not None and output_cost_per_token is not None + ) or input_cost_per_second is not None: litellm.register_model( { - f"{custom_llm_provider}/{model}": { - "input_cost_per_token": input_cost_per_token, - "output_cost_per_token": output_cost_per_token, - "litellm_provider": custom_llm_provider, - } - } - ) - elif ( - input_cost_per_second is not None - ): # time based pricing just needs cost in place - output_cost_per_second = output_cost_per_second - litellm.register_model( - { - f"{custom_llm_provider}/{model}": { - "input_cost_per_second": input_cost_per_second, - "output_cost_per_second": output_cost_per_second, - "litellm_provider": custom_llm_provider, - } + f"{custom_llm_provider}/{model}": _build_custom_pricing_entry( + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=model_info, + ) } ) ### BUILD CUSTOM PROMPT TEMPLATE -- IF GIVEN ### @@ -4644,7 +4660,6 @@ def embedding( # noqa: PLR0915 input_cost_per_token = kwargs.get("input_cost_per_token", None) output_cost_per_token = kwargs.get("output_cost_per_token", None) input_cost_per_second = kwargs.get("input_cost_per_second", None) - output_cost_per_second = kwargs.get("output_cost_per_second", None) openai_params = [ "user", "dimensions", @@ -4694,25 +4709,16 @@ def embedding( # noqa: PLR0915 ) ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### - if input_cost_per_token is not None and output_cost_per_token is not None: + if ( + input_cost_per_token is not None and output_cost_per_token is not None + ) or input_cost_per_second is not None: litellm.register_model( { - f"{custom_llm_provider}/{model}": { - "input_cost_per_token": input_cost_per_token, - "output_cost_per_token": output_cost_per_token, - "litellm_provider": custom_llm_provider, - } - } - ) - if input_cost_per_second is not None: # time based pricing just needs cost in place - output_cost_per_second = output_cost_per_second or 0.0 - litellm.register_model( - { - f"{custom_llm_provider}/{model}": { - "input_cost_per_second": input_cost_per_second, - "output_cost_per_second": output_cost_per_second, - "litellm_provider": custom_llm_provider, - } + f"{custom_llm_provider}/{model}": _build_custom_pricing_entry( + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=kwargs.get("model_info"), + ) } ) diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py new file mode 100644 index 00000000000..fa023d9943e --- /dev/null +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -0,0 +1,180 @@ +""" +Test that register_model() in completion() and embedding() passes all +custom pricing fields from kwargs and model_info, not just the base +input/output costs. + +Previously, only input_cost_per_token, output_cost_per_token, and +litellm_provider were forwarded. Fields like cache_read_input_token_cost, +mode, and supports_prompt_caching were dropped, causing incorrect cost +calculations for DB-sourced models with prompt caching pricing. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.main import _build_custom_pricing_entry + + +def test_build_custom_pricing_entry_includes_all_kwargs_fields(): + """All CustomPricingLiteLLMParams fields present in kwargs should be + included in the resulting entry dict.""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "cache_read_input_token_cost": 0.00025, + "cache_creation_input_token_cost": 0.005, + "output_cost_per_reasoning_token": 0.01, + "input_cost_per_audio_token": 0.003, + "unrelated_kwarg": "should_be_ignored", + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + ) + + assert entry["litellm_provider"] == "openai" + assert entry["input_cost_per_token"] == 0.001 + assert entry["output_cost_per_token"] == 0.002 + assert entry["cache_read_input_token_cost"] == 0.00025 + assert entry["cache_creation_input_token_cost"] == 0.005 + assert entry["output_cost_per_reasoning_token"] == 0.01 + assert entry["input_cost_per_audio_token"] == 0.003 + assert "unrelated_kwarg" not in entry + + +def test_build_custom_pricing_entry_merges_model_info_metadata(): + """Fields from model_info (mode, supports_prompt_caching, max_tokens) + should be merged into the entry when present.""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + model_info = { + "id": "deployment-123", + "mode": "chat", + "supports_prompt_caching": True, + "max_tokens": 128000, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + model_info=model_info, + ) + + assert entry["mode"] == "chat" + assert entry["supports_prompt_caching"] is True + assert entry["max_tokens"] == 128000 + + +def test_build_custom_pricing_entry_kwargs_take_precedence_over_model_info(): + """If a field appears in both kwargs and model_info, the kwargs value + should take precedence (setdefault behavior).""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + model_info = { + "mode": "chat", + "supports_prompt_caching": True, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + model_info=model_info, + ) + + # model_info fields should be set via setdefault + assert entry["mode"] == "chat" + assert entry["supports_prompt_caching"] is True + + +def test_build_custom_pricing_entry_skips_none_values(): + """Fields with None values in kwargs should not be included.""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": None, # explicitly None + "cache_read_input_token_cost": None, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + ) + + assert entry["input_cost_per_token"] == 0.001 + assert "output_cost_per_token" not in entry + assert "cache_read_input_token_cost" not in entry + + +def test_build_custom_pricing_entry_handles_no_model_info(): + """Should work correctly when model_info is None.""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + model_info=None, + ) + + assert entry["litellm_provider"] == "openai" + assert entry["input_cost_per_token"] == 0.001 + assert entry["output_cost_per_token"] == 0.002 + assert "mode" not in entry + + +def test_register_model_receives_cache_pricing_fields(): + """End-to-end: when register_model is called with a full pricing entry, + the cache pricing fields should be present in litellm.model_cost.""" + model_key = "openai/test-custom-model-with-cache-pricing" + + litellm.register_model( + { + model_key: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "cache_read_input_token_cost": 0.00025, + "supports_prompt_caching": True, + "mode": "chat", + "max_tokens": 8192, + "litellm_provider": "openai", + } + } + ) + + registered = litellm.model_cost.get(model_key) + assert registered is not None, f"{model_key} should be in model_cost" + assert registered["cache_read_input_token_cost"] == 0.00025 + assert registered["supports_prompt_caching"] is True + assert registered["mode"] == "chat" + assert registered["max_tokens"] == 8192 + + # Cleanup + litellm.model_cost.pop(model_key, None) + + +def test_build_custom_pricing_entry_time_based(): + """Time-based pricing fields should be included correctly.""" + kwargs = { + "input_cost_per_second": 0.01, + "output_cost_per_second": 0.02, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + ) + + assert entry["litellm_provider"] == "openai" + assert entry["input_cost_per_second"] == 0.01 + assert entry["output_cost_per_second"] == 0.02 From 58cd6f68995130d7a163bca3dcc0af088daf51c3 Mon Sep 17 00:00:00 2001 From: openhands Date: Mon, 2 Mar 2026 08:31:33 +0000 Subject: [PATCH 016/147] test: fix misleading precedence test per review feedback Renamed test_build_custom_pricing_entry_kwargs_take_precedence_over_model_info to test_build_custom_pricing_entry_setdefault_does_not_override_existing. The original test claimed to verify kwargs precedence over model_info but had no overlapping keys between the two sources. CustomPricingLiteLLMParams fields and the model_info metadata keys (mode, supports_prompt_caching, max_tokens) do not currently overlap. Updated the test to verify that model_info fields are correctly merged, and added an explicit setdefault assertion demonstrating that pre-existing keys would not be overwritten. Co-authored-by: openhands --- .../test_register_model_custom_pricing.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index fa023d9943e..1efd698fb64 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -73,9 +73,12 @@ def test_build_custom_pricing_entry_merges_model_info_metadata(): assert entry["max_tokens"] == 128000 -def test_build_custom_pricing_entry_kwargs_take_precedence_over_model_info(): - """If a field appears in both kwargs and model_info, the kwargs value - should take precedence (setdefault behavior).""" +def test_build_custom_pricing_entry_setdefault_does_not_override_existing(): + """model_info uses setdefault, so it should not override a key that is + already present in the entry dict. Currently CustomPricingLiteLLMParams + and the model_info keys (mode, supports_prompt_caching, max_tokens) do + not overlap, but if they ever do, setdefault ensures the kwargs-sourced + value wins.""" kwargs = { "input_cost_per_token": 0.001, "output_cost_per_token": 0.002, @@ -83,6 +86,7 @@ def test_build_custom_pricing_entry_kwargs_take_precedence_over_model_info(): model_info = { "mode": "chat", "supports_prompt_caching": True, + "max_tokens": 128000, } entry = _build_custom_pricing_entry( @@ -91,9 +95,17 @@ def test_build_custom_pricing_entry_kwargs_take_precedence_over_model_info(): model_info=model_info, ) - # model_info fields should be set via setdefault assert entry["mode"] == "chat" assert entry["supports_prompt_caching"] is True + assert entry["max_tokens"] == 128000 + + # Verify setdefault behavior: if a model_info key already exists in + # the entry (e.g. from a future CustomPricingLiteLLMParams addition), + # setdefault must not overwrite it. + entry["mode"] = "embedding" # simulate pre-existing value + # Re-apply setdefault the same way _build_custom_pricing_entry does + entry.setdefault("mode", model_info["mode"]) + assert entry["mode"] == "embedding" # must NOT revert to "chat" def test_build_custom_pricing_entry_skips_none_values(): From 538f11bdfa8c536424c0a043d73bc64c23d36835 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 16:53:53 +0530 Subject: [PATCH 017/147] feat(types): add _aresponses_websocket CallType and video_tokens to PromptTokensDetailsWrapper Co-Authored-By: Claude Sonnet 4.6 --- litellm/model_prices_and_context_window_backup.json | 2 +- litellm/types/utils.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f52288ea72a..cbd64a178b8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16289,7 +16289,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "supports_reasoning": false, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8b9359876e6..8d0cbdefd5b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -291,6 +291,7 @@ class CallTypes(str, Enum): search = "search" asearch = "asearch" arealtime = "_arealtime" + aresponses_websocket = "_aresponses_websocket" create_batch = "create_batch" acreate_batch = "acreate_batch" aretrieve_batch = "aretrieve_batch" @@ -1398,6 +1399,9 @@ class PromptTokensDetailsWrapper( image_tokens: Optional[int] = None """Image tokens sent to the model.""" + video_tokens: Optional[int] = None + """Video tokens sent to the model.""" + web_search_requests: Optional[int] = None """Number of web search requests made by the tool call. Used for Anthropic to calculate web search cost.""" From 19ce26501b30fc97b1be6a83a2848e7c5445c96e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 16:55:27 +0530 Subject: [PATCH 018/147] feat(responses): add WebSocket streaming iterator for responses API Co-Authored-By: Claude Sonnet 4.6 --- litellm/responses/streaming_iterator.py | 189 +++++++++++++++++++++++- 1 file changed, 187 insertions(+), 2 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 43ef4610b4b..b6b90067875 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -3,12 +3,15 @@ import json import time import traceback from datetime import datetime -from typing import Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional import httpx import litellm -from litellm.constants import LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING +from litellm.constants import ( + LITELLM_MAX_STREAMING_DURATION_SECONDS, + STREAM_SSE_DONE_STRING, +) from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -654,3 +657,185 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): for c in getattr(out_item, "content", []): out += c.text return out + + +# --------------------------------------------------------------------------- +# WebSocket mode streaming (bidirectional forwarding) +# --------------------------------------------------------------------------- + +if TYPE_CHECKING: + from websockets.asyncio.client import ClientConnection as _WsClientConnection + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.thread_pool_executor import executor as _ws_executor + +RESPONSES_WS_LOGGED_EVENT_TYPES = [ + "response.created", + "response.completed", + "response.failed", + "response.incomplete", + "error", +] + + +class ResponsesWebSocketStreaming: + """ + Manages bidirectional WebSocket forwarding for the Responses API + WebSocket mode (wss://.../v1/responses). + + Unlike the Realtime API, the Responses API WebSocket mode: + - Uses response.create as the client-to-server event + - Streams back the same events as the HTTP streaming Responses API + - Supports previous_response_id for incremental continuation + - Supports generate: false for warmup + - One response at a time per connection (sequential, no multiplexing) + """ + + def __init__( + self, + websocket: Any, + backend_ws: Any, + logging_obj: LiteLLMLoggingObj, + user_api_key_dict: Optional[Any] = None, + request_data: Optional[Dict] = None, + ): + self.websocket = websocket + self.backend_ws = backend_ws + self.logging_obj = logging_obj + self.user_api_key_dict = user_api_key_dict + self.request_data: Dict = request_data or {} + self.messages: list[Dict] = [] + self.input_messages: list[Dict[str, str]] = [] + + def _should_store_event(self, event_obj: dict) -> bool: + return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES + + def _store_event(self, event: Any) -> None: + if isinstance(event, bytes): + event = event.decode("utf-8") + if isinstance(event, str): + try: + event_obj = json.loads(event) + except (json.JSONDecodeError, TypeError): + return + else: + event_obj = event + + if self._should_store_event(event_obj): + self.messages.append(event_obj) + + def _collect_input_from_client_event(self, message: Any) -> None: + """Extract user input content from response.create for logging.""" + try: + if isinstance(message, str): + msg_obj = json.loads(message) + elif isinstance(message, dict): + msg_obj = message + else: + return + + if msg_obj.get("type") != "response.create": + return + + input_items = msg_obj.get("input", []) + if isinstance(input_items, str): + self.input_messages.append({"role": "user", "content": input_items}) + return + + if isinstance(input_items, list): + for item in input_items: + if not isinstance(item, dict): + continue + if item.get("type") == "message" and item.get("role") == "user": + content = item.get("content", []) + if isinstance(content, str): + self.input_messages.append( + {"role": "user", "content": content} + ) + elif isinstance(content, list): + for c in content: + if ( + isinstance(c, dict) + and c.get("type") == "input_text" + ): + text = c.get("text", "") + if text: + self.input_messages.append( + {"role": "user", "content": text} + ) + except (json.JSONDecodeError, AttributeError, TypeError): + pass + + def _store_input(self, message: Any) -> None: + self._collect_input_from_client_event(message) + if self.logging_obj: + self.logging_obj.pre_call(input=message, api_key="") + + async def _log_messages(self) -> None: + if not self.logging_obj: + return + if self.input_messages: + self.logging_obj.model_call_details["messages"] = self.input_messages + if self.messages: + asyncio.create_task( + self.logging_obj.async_success_handler(self.messages) + ) + _ws_executor.submit(self.logging_obj.success_handler, self.messages) + + async def backend_to_client(self) -> None: + """Forward events from backend WebSocket to the client.""" + import websockets + + try: + while True: + try: + raw_response = await self.backend_ws.recv(decode=False) # type: ignore[union-attr] + except TypeError: + raw_response = await self.backend_ws.recv() # type: ignore[union-attr, assignment] + + if isinstance(raw_response, bytes): + response_str = raw_response.decode("utf-8") + else: + response_str = raw_response + + self._store_event(response_str) + await self.websocket.send_text(response_str) + + except websockets.exceptions.ConnectionClosed as e: # type: ignore + verbose_logger.debug( + "Responses WS backend connection closed: %s", e + ) + except Exception as e: + verbose_logger.exception( + "Error in responses WS backend_to_client: %s", e + ) + finally: + await self._log_messages() + + async def client_to_backend(self) -> None: + """Forward response.create events from client to backend.""" + try: + while True: + message = await self.websocket.receive_text() + + self._store_input(message) + self._store_event(message) + await self.backend_ws.send(message) # type: ignore[union-attr] + + except Exception as e: + verbose_logger.debug("Responses WS client_to_backend ended: %s", e) + + async def bidirectional_forward(self) -> None: + """Run both forwarding directions concurrently.""" + forward_task = asyncio.create_task(self.backend_to_client()) + try: + await self.client_to_backend() + except Exception: + pass + finally: + if not forward_task.done(): + forward_task.cancel() + try: + await forward_task + except asyncio.CancelledError: + pass From 39748fd4a3327f6e7e0796fe53a49d9e14ab62b9 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 17:21:47 +0530 Subject: [PATCH 019/147] feat(responses): add _aresponses_websocket function and HTTP handler support for WebSocket mode Also fix pyrightconfig.json to use the conda venv for type checking, and remove redundant inline import of ResponsesAPIRequestUtils that was confusing pyright. Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/llm_http_handler.py | 93 ++++++++++++++++ litellm/responses/main.py | 101 +++++++++++++++++- pyrightconfig.json | 4 +- 3 files changed, 195 insertions(+), 3 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d6fdc58099f..29d494dbb50 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -69,6 +69,7 @@ from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, ResponsesAPIStreamingIterator, + ResponsesWebSocketStreaming, SyncResponsesAPIStreamingIterator, ) from litellm.types.containers.main import ( @@ -4731,6 +4732,98 @@ class BaseLLMHTTPHandler: f"Unexpected error while closing WebSocket: {close_error}" ) + async def async_responses_websocket( + self, + model: str, + websocket: Any, + logging_obj: LiteLLMLoggingObj, + responses_api_provider_config: BaseResponsesAPIConfig, + api_base: Optional[str] = None, + api_key: Optional[str] = None, + timeout: Optional[float] = None, + user_api_key_dict: Optional[Any] = None, + litellm_metadata: Optional[Dict[str, Any]] = None, + ): + """ + Handles Responses API WebSocket mode. + + Opens a persistent WebSocket to the provider's /v1/responses endpoint + and proxies response.create events bidirectionally for lower-latency + agentic workflows. + """ + import websockets + from websockets.asyncio.client import ClientConnection + + litellm_params = GenericLiteLLMParams() + headers = responses_api_provider_config.validate_environment( + headers={}, + model=model, + litellm_params=litellm_params, + ) + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + http_url = responses_api_provider_config.get_complete_url( + api_base=api_base, + litellm_params={}, + ) + # /responses -> wss:// URL + ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + + try: + ssl_context = get_shared_realtime_ssl_context() + if ws_url.startswith("wss://") and ssl_context is False: + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + + logging_obj.pre_call( + input=None, + api_key=api_key or "", + additional_args={ + "api_base": ws_url, + "headers": headers, + "complete_input_dict": {"mode": "responses_websocket"}, + }, + ) + + async with websockets.connect( # type: ignore + ws_url, + additional_headers=headers, + max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=ssl_context, + ) as backend_ws: + _request_data: Dict[str, Any] = {} + if litellm_metadata: + _request_data["litellm_metadata"] = litellm_metadata + streaming = ResponsesWebSocketStreaming( + websocket=websocket, + backend_ws=cast(ClientConnection, backend_ws), + logging_obj=logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=_request_data, + ) + await streaming.bidirectional_forward() + + except websockets.exceptions.InvalidStatusCode as e: # type: ignore + verbose_logger.exception(f"Error connecting to responses WS backend: {e}") + await websocket.close(code=e.status_code, reason=str(e)) + except Exception as e: + verbose_logger.exception(f"Error in responses WS: {e}") + try: + await websocket.close( + code=1011, reason=f"Internal server error: {str(e)}" + ) + except RuntimeError as close_error: + if "already completed" in str(close_error) or "websocket.close" in str( + close_error + ): + pass + else: + raise Exception( + f"Unexpected error while closing WebSocket: {close_error}" + ) + def image_edit_handler( self, model: str, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 05fd6026af2..6bdeaf66e62 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -51,6 +51,8 @@ if TYPE_CHECKING: from litellm.types.llms.openai import ResponseText # type: ignore else: ResponseText = str # Fallback for ResponseText import +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -182,8 +184,6 @@ async def aresponses_api_with_mcp( mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None secret_fields = kwargs.get("secret_fields") if secret_fields and isinstance(secret_fields, dict): - from litellm.responses.utils import ResponsesAPIRequestUtils - mcp_auth_header, mcp_server_auth_headers, _, _ = ( ResponsesAPIRequestUtils.extract_mcp_headers_from_request( secret_fields=secret_fields, tools=tools @@ -1651,3 +1651,100 @@ def compact_responses( completion_kwargs=local_vars, extra_kwargs=kwargs, ) + + +# --------------------------------------------------------------------------- +# Responses API WebSocket mode +# --------------------------------------------------------------------------- + + +def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: + metadata: dict = {**(kwargs.get("litellm_metadata") or {})} + guardrails = ( + (kwargs.get("metadata") or {}).get("guardrails") + or kwargs.get("guardrails") + or [] + ) + if guardrails: + metadata["guardrails"] = guardrails + return metadata + + +@client +async def _aresponses_websocket( + model: str, + websocket: Any, + api_base: Optional[str] = None, + api_key: Optional[str] = None, + timeout: Optional[float] = None, + **kwargs, +): + """ + Private function to handle the Responses API WebSocket mode. + + For PROXY use only. + + Resolves the LLM provider from ``model``, looks up the matching + ``BaseResponsesAPIConfig``, and hands off to + ``BaseLLMHTTPHandler.async_responses_websocket``. + """ + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + user = kwargs.get("user", None) + litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params_dict = get_litellm_params(**kwargs) + + model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = ( + litellm.get_llm_provider( + model=model, + api_base=api_base, + api_key=api_key, + ) + ) + + litellm_logging_obj.update_environment_variables( + model=model, + user=user, + optional_params={}, + litellm_params=litellm_params_dict, + custom_llm_provider=_custom_llm_provider, + ) + + responses_api_provider_config: Optional[BaseResponsesAPIConfig] = None + if _custom_llm_provider is not None: + responses_api_provider_config = ( + ProviderConfigManager.get_provider_responses_api_config( + model=model, + provider=litellm.LlmProviders(_custom_llm_provider), + ) + ) + + if responses_api_provider_config is None: + raise ValueError( + f"Responses API WebSocket mode is not supported for provider: {_custom_llm_provider}" + ) + + resolved_api_base = ( + dynamic_api_base + or litellm_params.api_base + or litellm.api_base + or None + ) + resolved_api_key = ( + dynamic_api_key + or litellm_params.api_key + or litellm.api_key + or litellm.openai_key + or get_secret_str("OPENAI_API_KEY") + ) + + await base_llm_http_handler.async_responses_websocket( + model=model, + websocket=websocket, + logging_obj=litellm_logging_obj, + responses_api_provider_config=responses_api_provider_config, + api_base=resolved_api_base, + api_key=resolved_api_key, + timeout=timeout, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata_for_ws(kwargs), + ) diff --git a/pyrightconfig.json b/pyrightconfig.json index f930e44d305..ec0a1823038 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -2,6 +2,8 @@ "ignore": [], "exclude": ["**/node_modules", "**/__pycache__", "litellm/types/utils.py", "litellm/proxy/_types.py"], "reportMissingImports": false, - "reportPrivateImportUsage": false + "reportPrivateImportUsage": false, + "venvPath": "/Users/sameerkankute/miniconda3/envs", + "venv": "litellm-dev" } \ No newline at end of file From eeb2d28621e5c8ed0bacab2934d63249e50eb9c2 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 17:22:00 +0530 Subject: [PATCH 020/147] feat(proxy): add _aresponses_websocket to common_request_processing route_type Literal Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/common_request_processing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1269f58213a..0eb6fc18486 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -518,6 +518,7 @@ class ProxyBaseLLMRequestProcessing: "aembedding", "aresponses", "_arealtime", + "_aresponses_websocket", "aget_responses", "adelete_responses", "acancel_responses", From 512a4389356e3f0f9bf2c98e9b3566000e5bd606 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 17:22:44 +0530 Subject: [PATCH 021/147] feat(proxy): add WebSocket endpoint for responses API and route_llm_request support Co-Authored-By: Claude Sonnet 4.6 --- .../proxy/response_api_endpoints/endpoints.py | 129 +++++++++++++++++- litellm/proxy/route_llm_request.py | 2 + 2 files changed, 129 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 44e8c42b2c1..4253c2ca832 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,14 +1,21 @@ import asyncio +import json import time -from typing import Any, AsyncIterator, Optional, cast +from typing import Any, AsyncIterator, Dict, Optional, cast from uuid import uuid4 +import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response +from starlette.websockets import WebSocket from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy._types import * -from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth +from litellm.proxy.auth.user_api_key_auth import ( + UserAPIKeyAuth, + user_api_key_auth, + user_api_key_auth_websocket, +) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.responses.main import DeleteResponseResult @@ -904,3 +911,121 @@ async def cancel_response( proxy_logging_obj=proxy_logging_obj, version=version, ) + + +@router.websocket("/v1/responses") +@router.websocket("/responses") +async def responses_websocket_endpoint( + websocket: WebSocket, + model: str = fastapi.Query( + ..., description="The model to use for the responses WebSocket session." + ), + user_api_key_dict=Depends(user_api_key_auth_websocket), +): + """ + Responses API WebSocket mode endpoint. + + Keeps a persistent WebSocket connection for response.create events, + enabling lower-latency agentic workflows with many tool-call round trips. + + See: https://developers.openai.com/api/docs/guides/websocket-mode/ + """ + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + from litellm.proxy.route_llm_request import route_request + + # Accept the WebSocket handshake + requested_protocols = [ + p.strip() + for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") + if p.strip() + ] + accept_kwargs: dict = {} + if requested_protocols: + accept_kwargs["subprotocol"] = requested_protocols[0] + await websocket.accept(**accept_kwargs) + + data: Dict[str, Any] = { + "model": model, + "websocket": websocket, + } + + # Construct a synthetic Request for pre-call processing + headers_list = list(websocket.scope.get("headers") or []) + scope: Dict[str, Any] = { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": headers_list, + } + request = Request(scope=scope) + request._url = websocket.url + + async def return_body(): + return f'{{"model": "{model}"}}'.encode() + + request.body = return_body # type: ignore + + # Phase 1: pre-call processing (auth, guardrails, rate limits) + base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) + try: + ( + data, + litellm_logging_obj, + ) = await base_llm_response_processor.common_processing_pre_call_logic( + request=request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + model=model, + route_type="_aresponses_websocket", + ) + except Exception as e: + verbose_proxy_logger.exception("Responses WebSocket pre-call error") + try: + await websocket.send_text( + json.dumps( + { + "type": "error", + "error": { + "type": "pre_call_error", + "message": str(e), + }, + } + ) + ) + except Exception: + pass + await websocket.close(code=1011, reason="Pre-call error") + return + + # Phase 2: route to upstream provider + try: + data["user_api_key_dict"] = user_api_key_dict + llm_call = await route_request( + data=data, + route_type="_aresponses_websocket", + llm_router=llm_router, + user_model=user_model, + ) + await llm_call + except Exception: + verbose_proxy_logger.exception("Responses WebSocket error") + await websocket.close(code=1011, reason="Internal server error") diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 63bd67abea2..1b791980af3 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -42,6 +42,7 @@ ROUTE_ENDPOINT_MAPPING = { "amoderation": "/moderations", "arerank": "/rerank", "aresponses": "/responses", + "_aresponses_websocket": "/responses", "alist_input_items": "/responses/{response_id}/input_items", "aimage_edit": "/images/edits", "acancel_responses": "/responses/{response_id}/cancel", @@ -163,6 +164,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin "acreate_response_reply", "alist_input_items", "_arealtime", # private function for realtime API + "_aresponses_websocket", # private function for responses WebSocket mode "aimage_edit", "agenerate_content", "agenerate_content_stream", From 0921e26b8c237e09c2fa607fc782adb77f1d3071 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 17:24:12 +0530 Subject: [PATCH 022/147] feat(router): register _aresponses_websocket in Router factory functions Also suppress pre-existing pyright type error on model_response.close() call which is guarded by hasattr at runtime. Co-Authored-By: Claude Sonnet 4.6 --- litellm/router.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 5c8d27e76d7..8f8c07408be 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -882,6 +882,9 @@ class Router: self._arealtime = self.factory_function( litellm._arealtime, call_type="_arealtime" ) + self._aresponses_websocket = self.factory_function( + litellm._aresponses_websocket, call_type="_aresponses_websocket" + ) self.acreate_fine_tuning_job = self.factory_function( litellm.acreate_fine_tuning_job, call_type="acreate_fine_tuning_job" ) @@ -1824,7 +1827,7 @@ class Router: finally: if hasattr(model_response, "close"): try: - model_response.close() + model_response.close() # type: ignore[reportAttributeAccessIssue] except BaseException as close_err: verbose_router_logger.debug( "stream_with_fallbacks: error closing model_response: %s", @@ -4659,6 +4662,7 @@ class Router: "afile_delete", "afile_content", "_arealtime", + "_aresponses_websocket", "acreate_fine_tuning_job", "acancel_fine_tuning_job", "alist_fine_tuning_jobs", @@ -4831,6 +4835,7 @@ class Router: "anthropic_messages", "aresponses", "_arealtime", + "_aresponses_websocket", "acreate_fine_tuning_job", "acancel_fine_tuning_job", "alist_fine_tuning_jobs", From 76ddfa184cbaac264af246c23ba44758165b9cba Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 17:24:27 +0530 Subject: [PATCH 023/147] feat(init): export _aresponses_websocket from litellm package Co-Authored-By: Claude Sonnet 4.6 --- litellm/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/__init__.py b/litellm/__init__.py index 2522d190570..aacd4299f19 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1246,6 +1246,7 @@ from .ocr.main import * from .rag.main import * from .search.main import * from .realtime_api.main import _arealtime +from .responses.main import _aresponses_websocket from .fine_tuning.main import * from .files.main import * from .vector_store_files.main import ( From 82f5055d89d31287d1db03b8de0f8148151a012e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 17:24:39 +0530 Subject: [PATCH 024/147] test(responses): add end-to-end test for responses API WebSocket mode Co-Authored-By: Claude Sonnet 4.6 --- .../test_responses_websocket_proxy_e2e.py | 239 ++++++++++++++++++ 1 file changed, 239 insertions(+) create mode 100644 tests/openai_endpoints_tests/test_responses_websocket_proxy_e2e.py diff --git a/tests/openai_endpoints_tests/test_responses_websocket_proxy_e2e.py b/tests/openai_endpoints_tests/test_responses_websocket_proxy_e2e.py new file mode 100644 index 00000000000..e76135baa7e --- /dev/null +++ b/tests/openai_endpoints_tests/test_responses_websocket_proxy_e2e.py @@ -0,0 +1,239 @@ +""" +E2E tests for OpenAI Responses API WebSocket mode through the LiteLLM proxy. + +Connects to ws://0.0.0.0:4000/v1/responses, sends response.create events, +and validates the streamed response events. + +Requires: + - Proxy running: python -m litellm.proxy.proxy_cli --config --port 4000 + - Model configured in proxy (e.g. gpt-4o-mini) + +See: https://developers.openai.com/api/docs/guides/websocket-mode/ +""" + +import asyncio +import json +import os + +import httpx +import pytest + +# ── Configuration ───────────────────────────────────────────────────────────── +PROXY_BASE_URL = os.environ.get("LITELLM_PROXY_BASE_URL", "ws://0.0.0.0:4000") +PROXY_MASTER_KEY = os.environ.get("LITELLM_PROXY_KEY", "sk-1234") +PROXY_MODEL = os.environ.get("LITELLM_PROXY_RESPONSES_MODEL", "gpt-4o-mini") +# ────────────────────────────────────────────────────────────────────────────── + + +def _generate_key() -> str: + """Generate a key for testing via proxy key/generate endpoint.""" + url = "http://0.0.0.0:4000/key/generate" + headers = { + "Authorization": f"Bearer {PROXY_MASTER_KEY}", + "Content-Type": "application/json", + } + response = httpx.post(url, headers=headers, json={}, timeout=10) + if response.status_code != 200: + raise Exception( + f"Key generation failed with status: {response.status_code}. " + "Is the proxy running?" + ) + return response.json()["key"] + + +def _assert_basic_response(events: list[dict], label: str = "") -> None: + """Assert that events contain response.created, response.completed, and usage.""" + prefix = f"[{label}] " if label else "" + types = [e.get("type") for e in events] + assert len(events) > 0, f"{prefix}no events received" + assert "response.created" in types, f"{prefix}missing response.created, got: {types}" + assert "response.completed" in types, ( + f"{prefix}missing response.completed, got: {types}" + ) + completed = next(e for e in events if e.get("type") == "response.completed") + resp = completed.get("response", {}) + assert resp.get("status") == "completed", ( + f"{prefix}status != completed: {resp.get('status')}" + ) + usage = resp.get("usage", {}) + assert usage.get("input_tokens", 0) > 0, f"{prefix}input_tokens=0" + assert usage.get("output_tokens", 0) > 0, f"{prefix}output_tokens=0" + streaming_types = { + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_item.done", + } + found = streaming_types & set(types) + assert found, f"{prefix}no streaming delta events found, got: {types}" + + +@pytest.mark.asyncio +async def test_responses_websocket_proxy_basic(): + """ + Sends a simple response.create event to the proxy WebSocket endpoint + and validates response.created, response.completed, and streaming events. + """ + try: + import websockets + except ImportError: + pytest.skip("websockets not installed") + + try: + key = _generate_key() + except Exception as e: + pytest.skip( + f"Proxy not available or key generation failed: {e}. " + "Start proxy: python -m litellm.proxy.proxy_cli --config --port 4000" + ) + + url = f"{PROXY_BASE_URL}/v1/responses?model={PROXY_MODEL}" + headers = {"Authorization": f"Bearer {key}"} + events: list[dict] = [] + + try: + async with websockets.connect( + url, additional_headers=headers, open_timeout=5 + ) as ws: + payload = { + "type": "response.create", + "model": PROXY_MODEL, + "store": False, + "input": [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Say hello in one word."} + ], + } + ], + "tools": [], + } + await ws.send(json.dumps(payload)) + for _ in range(50): + msg = await asyncio.wait_for(ws.recv(), timeout=15) + event = json.loads(msg) + events.append(event) + if event.get("type") in ( + "response.completed", + "response.failed", + "error", + ): + break + except Exception as e: + pytest.fail( + f"WebSocket connection failed: {e}. " + "Ensure proxy is running and model is configured." + ) + + _assert_basic_response(events, "proxy-basic") + + +@pytest.mark.asyncio +async def test_responses_websocket_proxy_multi_turn(): + """ + Sends two sequential response.create events with previous_response_id + to validate multi-turn conversation over a single WebSocket. + """ + try: + import websockets + except ImportError: + pytest.skip("websockets not installed") + + try: + key = _generate_key() + except Exception as e: + pytest.skip( + f"Proxy not available or key generation failed: {e}. " + "Start proxy: python -m litellm.proxy.proxy_cli --config --port 4000" + ) + + url = f"{PROXY_BASE_URL}/v1/responses?model={PROXY_MODEL}" + headers = {"Authorization": f"Bearer {key}"} + all_events: list[dict] = [] + completed: list[dict] = [] + first_id = None + + try: + async with websockets.connect( + url, additional_headers=headers, open_timeout=5 + ) as ws: + # Turn 1 + await ws.send( + json.dumps( + { + "type": "response.create", + "model": PROXY_MODEL, + "store": True, + "input": [ + { + "type": "message", + "role": "user", + "content": [ + { + "type": "input_text", + "text": "Remember the number 7. Just say OK.", + } + ], + } + ], + } + ) + ) + for _ in range(50): + msg = await asyncio.wait_for(ws.recv(), timeout=15) + event = json.loads(msg) + all_events.append(event) + if event.get("type") == "response.completed": + completed.append(event) + first_id = event.get("response", {}).get("id") + break + if event.get("type") in ("response.failed", "error"): + break + + assert first_id, "Turn 1 never completed" + + # Turn 2 + await ws.send( + json.dumps( + { + "type": "response.create", + "model": PROXY_MODEL, + "store": True, + "previous_response_id": first_id, + "input": [ + { + "type": "message", + "role": "user", + "content": [ + { + "type": "input_text", + "text": "What number did I tell you to remember?", + } + ], + } + ], + } + ) + ) + for _ in range(50): + msg = await asyncio.wait_for(ws.recv(), timeout=15) + event = json.loads(msg) + all_events.append(event) + if event.get("type") == "response.completed": + completed.append(event) + break + if event.get("type") in ("response.failed", "error"): + break + + except Exception as e: + pytest.fail( + f"WebSocket multi-turn failed: {e}. " + "Ensure proxy is running and model is configured." + ) + + assert len(completed) >= 2, ( + f"Expected 2 response.completed events, got {len(completed)}" + ) + assert completed[1].get("response", {}).get("status") == "completed" From c1136348f050170a686399bf0c4bf33027c91171 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 17:27:24 +0530 Subject: [PATCH 025/147] revert pyrightconfig --- pyrightconfig.json | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pyrightconfig.json b/pyrightconfig.json index ec0a1823038..f930e44d305 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -2,8 +2,6 @@ "ignore": [], "exclude": ["**/node_modules", "**/__pycache__", "litellm/types/utils.py", "litellm/proxy/_types.py"], "reportMissingImports": false, - "reportPrivateImportUsage": false, - "venvPath": "/Users/sameerkankute/miniconda3/envs", - "venv": "litellm-dev" + "reportPrivateImportUsage": false } \ No newline at end of file From cc650f4865a692997714fd81b5bc0324d6c17f32 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 18:30:39 +0530 Subject: [PATCH 026/147] fix(responses): add in-memory session tracking to ManagedResponsesWebSocketHandler for previous_response_id Spend logs are written asynchronously in batches, so a DB lookup for a just-completed response's spend log races and returns zero rows. Replace the DB-only fallback with an in-memory session store (_session_history) that is populated after each response.completed event and consulted at the start of the next response.create. This makes same-connection multi-turn reliable without any timing dependency on the DB write queue. Made-with: Cursor --- litellm/responses/streaming_iterator.py | 411 +++++++++++++++++++++++- ruff.toml | 1 + 2 files changed, 411 insertions(+), 1 deletion(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index b6b90067875..f026ce46af5 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -3,7 +3,7 @@ import json import time import traceback from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional import httpx @@ -839,3 +839,412 @@ class ResponsesWebSocketStreaming: await forward_task except asyncio.CancelledError: pass + try: + await self.backend_ws.close() + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Managed WebSocket mode (HTTP-backed, provider-agnostic) +# --------------------------------------------------------------------------- + +_RESPONSE_CREATE_PARAMS = ( + "input", + "model", + "previous_response_id", + "instructions", + "max_output_tokens", + "tools", + "tool_choice", + "temperature", + "top_p", + "store", + "metadata", + "truncation", + "reasoning", + "stream", + "include", + "parallel_tool_calls", + "text", + "user", + "service_tier", + "safety_identifier", + "background", +) + +_MANAGED_WS_SKIP_KWARGS = frozenset( + { + "litellm_logging_obj", + "litellm_call_id", + "aresponses", + "_aresponses_websocket", + "user_api_key_dict", + } +) + + +class ManagedResponsesWebSocketHandler: + """ + Handles Responses API WebSocket mode for providers that do not expose a + native ``wss://`` responses endpoint. + + Instead of proxying to a provider WebSocket, this handler: + - Listens for ``response.create`` events from the client + - Makes HTTP streaming calls via ``litellm.aresponses(stream=True)`` + - Serialises and forwards every streaming event back over the WebSocket + - Supports ``previous_response_id`` for multi-turn conversations via + in-memory session tracking (avoids async DB-write timing issues) + - Supports sequential requests over a single persistent connection + + This makes every provider that LiteLLM can reach over HTTP available on + the WebSocket transport without any provider-specific changes. + """ + + def __init__( + self, + websocket: Any, + model: str, + logging_obj: "LiteLLMLoggingObj", + user_api_key_dict: Optional[Any] = None, + litellm_metadata: Optional[Dict[str, Any]] = None, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + timeout: Optional[float] = None, + custom_llm_provider: Optional[str] = None, + **kwargs: Any, + ) -> None: + self.websocket = websocket + self.model = model + self.logging_obj = logging_obj + self.user_api_key_dict = user_api_key_dict + self.litellm_metadata: Dict[str, Any] = litellm_metadata or {} + self.api_key = api_key + self.api_base = api_base + self.timeout = timeout + self.custom_llm_provider = custom_llm_provider + # Carry through safe pass-through kwargs (e.g. extra_headers) + self.extra_kwargs: Dict[str, Any] = { + k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS + } + # In-memory session history: response_id → list of input+output messages. + # Keyed by the DECODED (pre-encoding) response ID from response.completed. + # This avoids the async DB-write race condition where spend logs haven't + # been committed yet when the next response.create arrives. + self._session_history: Dict[str, List[Dict[str, Any]]] = {} + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _serialize_chunk(chunk: Any) -> Optional[str]: + """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" + try: + if hasattr(chunk, "model_dump_json"): + return chunk.model_dump_json(exclude_none=True) + if hasattr(chunk, "model_dump"): + return json.dumps(chunk.model_dump(exclude_none=True), default=str) + if isinstance(chunk, dict): + return json.dumps(chunk, default=str) + return json.dumps(str(chunk)) + except Exception as exc: + verbose_logger.debug("ManagedResponsesWS: failed to serialize chunk: %s", exc) + return None + + async def _send_error(self, message: str, error_type: str = "server_error") -> None: + try: + await self.websocket.send_text( + json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) + ) + except Exception: + pass + + # ------------------------------------------------------------------ + # Core request handler + # ------------------------------------------------------------------ + + def _get_history_messages(self, previous_response_id: str) -> List[Dict[str, Any]]: + """ + Return accumulated message history for *previous_response_id*. + + Checks the in-memory session store first (fast path, no DB round-trip). + The key is the *decoded* response ID (the raw provider response ID before + LiteLLM base64-encodes it into the ``resp_...`` format). + """ + from litellm.responses.utils import ResponsesAPIRequestUtils + + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id( + previous_response_id + ) + raw_id = decoded.get("response_id", previous_response_id) + return list(self._session_history.get(raw_id, [])) + + def _store_history( + self, + response_id: str, + input_messages: List[Dict[str, Any]], + output_messages: List[Dict[str, Any]], + ) -> None: + """ + Persist a turn's messages in the in-memory session store. + + *response_id* is the raw (decoded) provider ID extracted from the + ``response.completed`` event so that the next turn can look it up via + :meth:`_get_history_messages`. + """ + prior: List[Dict[str, Any]] = self._session_history.get(response_id, []) + self._session_history[response_id] = prior + input_messages + output_messages + + @staticmethod + def _extract_response_id(completed_event: Dict[str, Any]) -> Optional[str]: + """ + Pull the raw (decoded) response ID out of a ``response.completed`` event. + Returns *None* if the event doesn't contain a usable ID. + """ + from litellm.responses.utils import ResponsesAPIRequestUtils + + resp_obj = completed_event.get("response", {}) + encoded_id: Optional[str] = resp_obj.get("id") if isinstance(resp_obj, dict) else None + if not encoded_id: + return None + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id) + return decoded.get("response_id", encoded_id) + + @staticmethod + def _extract_output_messages(completed_event: Dict[str, Any]) -> List[Dict[str, Any]]: + """ + Convert the output items in a ``response.completed`` event into + chat-completion style messages suitable for the next turn's ``input``. + """ + resp_obj = completed_event.get("response", {}) + if not isinstance(resp_obj, dict): + return [] + messages: List[Dict[str, Any]] = [] + for item in resp_obj.get("output", []) or []: + if not isinstance(item, dict): + continue + item_type = item.get("type") + role = item.get("role", "assistant") + if item_type == "message": + content_parts = item.get("content") or [] + text_parts = [ + p.get("text", "") + for p in content_parts + if isinstance(p, dict) and p.get("type") in ("output_text", "text") + ] + text = "".join(text_parts) + if text: + messages.append({"type": "message", "role": role, "content": [{"type": "output_text", "text": text}]}) + elif item_type == "function_call": + messages.append(item) + return messages + + @staticmethod + def _input_to_messages(input_val: Any) -> List[Dict[str, Any]]: + """ + Normalise the ``input`` field of a ``response.create`` event to a list + of Responses API message dicts. + """ + if isinstance(input_val, str): + return [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": input_val}]}] + if isinstance(input_val, list): + return [item for item in input_val if isinstance(item, dict)] + return [] + + async def _process_response_create(self, raw_message: str) -> None: + """ + Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``, + and forward every streaming event to the client. + + Multi-turn support via in-memory session history + ------------------------------------------------ + When ``previous_response_id`` is present in the event: + 1. Look up the accumulated message history in ``self._session_history`` + (keyed by the decoded provider response ID). + 2. Prepend those messages to the current ``input`` so the model has full + conversation context. + 3. After the stream completes, extract the new response ID and output + messages from ``response.completed`` and store them in + ``self._session_history`` for the next turn. + + This in-memory approach avoids the async DB-write race condition that + occurs when spend logs haven't been committed by the time the second + ``response.create`` arrives over the same WebSocket connection. + """ + import litellm as _litellm + + try: + msg_obj = json.loads(raw_message) + except json.JSONDecodeError: + await self._send_error("Invalid JSON in response.create event", "invalid_request_error") + return + + if msg_obj.get("type") != "response.create": + # Silently ignore non-response.create messages (e.g. warmup pings) + return + + # Support two wire formats: + # Nested : {"type": "response.create", "response": {"input": [...], ...}} + # Flat : {"type": "response.create", "input": [...], "model": "...", ...} + nested = msg_obj.get("response") + if isinstance(nested, dict) and nested: + response_params: Dict[str, Any] = nested + else: + response_params = {k: v for k, v in msg_obj.items() if k != "type"} + + # Build kwargs for aresponses from the response.create payload + call_kwargs: Dict[str, Any] = {} + for param in _RESPONSE_CREATE_PARAMS: + if param in response_params and response_params[param] is not None: + call_kwargs[param] = response_params[param] + + # Always stream + call_kwargs["stream"] = True + + # Use the model from the event if provided, otherwise fall back to the + # model supplied at WebSocket connect time. + event_model = call_kwargs.pop("model", None) + model = event_model or self.model + + # ---- In-memory multi-turn: prepend history when previous_response_id set ---- + previous_response_id: Optional[str] = call_kwargs.pop("previous_response_id", None) + current_input = call_kwargs.get("input") + current_messages = self._input_to_messages(current_input) + if previous_response_id: + history = self._get_history_messages(previous_response_id) + if history: + # Prepend history; current messages are the new user turn + call_kwargs["input"] = history + current_messages + verbose_logger.debug( + "ManagedResponsesWS: prepended %d history messages for previous_response_id=%s", + len(history), + previous_response_id, + ) + else: + verbose_logger.debug( + "ManagedResponsesWS: no in-memory history for previous_response_id=%s; " + "falling back to DB-based session reconstruction", + previous_response_id, + ) + # Fall back to DB-based session reconstruction (may work for + # cross-connection multi-turn when spend logs are committed) + call_kwargs["previous_response_id"] = previous_response_id + # --------------------------------------------------------------------------- + + # Inject connection-level credentials and metadata. + # Only propagate custom_llm_provider when the request is using the + # same model as the WebSocket connection (i.e. no per-request model + # override). If the payload specifies a different model, let litellm + # re-resolve the provider from the model name so we don't accidentally + # force the wrong backend. + if self.api_key is not None: + call_kwargs["api_key"] = self.api_key + if self.api_base is not None: + call_kwargs["api_base"] = self.api_base + if self.timeout is not None: + call_kwargs["timeout"] = self.timeout + if self.custom_llm_provider is not None and not event_model: + call_kwargs["custom_llm_provider"] = self.custom_llm_provider + if self.litellm_metadata: + call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) + + # Update proxy_server_request body so spend logs record the full request. + proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get( + "proxy_server_request" + ) or {} + if isinstance(proxy_server_request, dict): + body = dict(proxy_server_request.get("body") or {}) + body["input"] = call_kwargs.get("input") + body["store"] = call_kwargs.get("store") + body["model"] = model + for k in ("tools", "tool_choice", "instructions", "metadata"): + if k in call_kwargs and call_kwargs[k] is not None: + body[k] = call_kwargs[k] + proxy_server_request = dict(proxy_server_request) + proxy_server_request["body"] = body + if "litellm_metadata" not in call_kwargs: + call_kwargs["litellm_metadata"] = {} + call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request + call_kwargs.setdefault("litellm_params", {}) + call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request + + # Merge any safe pass-through kwargs (extra_headers, etc.) + call_kwargs.update(self.extra_kwargs) + + # Track the completed event to update in-memory history after the turn. + completed_event: Optional[Dict[str, Any]] = None + + try: + stream_response = await _litellm.aresponses(model=model, **call_kwargs) + + async for chunk in stream_response: # type: ignore[union-attr] + if chunk is None: + continue + serialized = self._serialize_chunk(chunk) + if serialized is not None: + # Capture the completed event for history bookkeeping + try: + chunk_dict = json.loads(serialized) if isinstance(serialized, str) else {} + if chunk_dict.get("type") == "response.completed": + completed_event = chunk_dict + except Exception: + pass + try: + await self.websocket.send_text(serialized) + except Exception as send_exc: + verbose_logger.debug( + "ManagedResponsesWS: error sending chunk to client: %s", send_exc + ) + return # Client disconnected + + except Exception as exc: + verbose_logger.exception("ManagedResponsesWS: error processing response.create: %s", exc) + await self._send_error(str(exc)) + return + + # ---- Store this turn in in-memory history for future previous_response_id lookups ---- + if completed_event is not None: + new_response_id = self._extract_response_id(completed_event) + if new_response_id: + output_msgs = self._extract_output_messages(completed_event) + # Accumulate: history from previous turn + current input + new output + prior_history: List[Dict[str, Any]] = [] + if previous_response_id: + prior_history = self._get_history_messages(previous_response_id) + self._store_history( + new_response_id, + prior_history + current_messages, + output_msgs, + ) + verbose_logger.debug( + "ManagedResponsesWS: stored %d messages for response_id=%s", + len(prior_history) + len(current_messages) + len(output_msgs), + new_response_id, + ) + # --------------------------------------------------------------------------- + + # ------------------------------------------------------------------ + # Main entry point + # ------------------------------------------------------------------ + + async def run(self) -> None: + """ + Main loop: accept ``response.create`` events sequentially and handle + each one before waiting for the next message. + """ + try: + while True: + try: + message = await self.websocket.receive_text() + except Exception as exc: + verbose_logger.debug( + "ManagedResponsesWS: client disconnected: %s", exc + ) + break + + await self._process_response_create(message) + + except Exception as exc: + verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc) + await self._send_error(f"Internal server error: {exc}") diff --git a/ruff.toml b/ruff.toml index 43ff802a684..76acb5dc936 100644 --- a/ruff.toml +++ b/ruff.toml @@ -16,3 +16,4 @@ exclude = ["litellm/types/*", "litellm/__init__.py", "litellm/proxy/example_conf "litellm/proxy/utils.py" = ["F401", "PLR0915"] "litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py" = ["PLR0915"] "litellm/proxy/guardrails/guardrail_hooks/guardrail_benchmarks/test_eval.py" = ["PLR0915"] +"litellm/responses/streaming_iterator.py" = ["PLR0915"] From e5037cd6eac6518dabc864fd7a390f0811a27ed7 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 2 Mar 2026 15:31:11 -0300 Subject: [PATCH 027/147] fix(vertex): preserve any-type schema semantics for JsonValue/Any fields Empty JSON schemas `{}` mean "any JSON value is valid" per the spec, but _build_vertex_schema was coercing them to `{"type": "object"}` in three places (process_items, convert_anyof_null_to_nullable, add_object_type), breaking Pydantic JsonValue fields on Gemini. Adds _is_any_type_schema() to detect unconstrained schemas and skip the type coercion, preserving Gemini's TYPE_UNSPECIFIED semantics. Fixes #22391 --- litellm/llms/vertex_ai/common_utils.py | 38 ++++++- .../vertex_ai/test_vertex_ai_common_utils.py | 107 ++++++++++++++++-- 2 files changed, 131 insertions(+), 14 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 02b69b94d94..8f8baad94c8 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -583,14 +583,40 @@ def _filter_anyof_fields(schema_dict: Dict[str, Any]) -> Dict[str, Any]: return schema_dict +def _is_any_type_schema(schema: dict) -> bool: + """ + Detect schemas that represent "any JSON value" (no type constraints). + + In JSON Schema, an empty schema {} means "any value is valid". + Schemas with only metadata keys (title, description, default, examples) + but no type-constraining keywords also represent "any type". + + Gemini's Schema proto uses TYPE_UNSPECIFIED (0) as default, + so omitting the type field is valid and means "any type". + """ + type_constraining_keys = { + "type", + "properties", + "items", + "anyOf", + "oneOf", + "allOf", + "enum", + "required", + "$ref", + "$schema", + } + return not any(key in type_constraining_keys for key in schema.keys()) + + def process_items(schema, depth=0): if depth > DEFAULT_MAX_RECURSE_DEPTH: raise ValueError( f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting." ) if isinstance(schema, dict): - if "items" in schema and schema["items"] == {}: - schema["items"] = {"type": "object"} + if "items" in schema and isinstance(schema["items"], dict) and _is_any_type_schema(schema["items"]): + pass # preserve "any type" semantics — don't coerce to object for key, value in schema.items(): if isinstance(value, dict): process_items(value, depth + 1) @@ -689,9 +715,8 @@ def convert_anyof_null_to_nullable(schema, depth=0): # remove null type anyof.remove(atype) contains_null = True - elif "type" not in atype and len(atype) == 0: - # Handle empty object case - atype["type"] = "object" + elif isinstance(atype, dict) and _is_any_type_schema(atype): + pass # preserve "any type" semantics — don't coerce to object if len(anyof) == 0: # Edge case: response schema with only null type present is invalid in Vertex AI @@ -726,7 +751,8 @@ def add_object_type(schema): # Gemini requires all function parameters to be type OBJECT # Handle case where schema has no properties and no type (e.g. tools with no arguments) if "type" not in schema and "anyOf" not in schema and "oneOf" not in schema and "allOf" not in schema: - schema["type"] = "object" + if not _is_any_type_schema(schema): + schema["type"] = "object" properties = schema.get("properties", None) if properties is not None: diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 94323e06901..b80aa996cae 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -212,7 +212,7 @@ def test_build_vertex_schema(): "properties": { "state": { "properties": { - "messages": {"items": {"type": "object"}, "type": "array"}, + "messages": {"items": {}, "type": "array"}, "conversation_id": {"type": "string"}, }, "required": ["messages", "conversation_id"], @@ -226,7 +226,7 @@ def test_build_vertex_schema(): "callbacks": { "anyOf": [ {"type": "array", "nullable": True}, - {"type": "object", "nullable": True}, + {"nullable": True}, ] }, "run_name": {"type": "string"}, @@ -270,23 +270,28 @@ def test_process_items_basic(): """Test basic functionality of process_items.""" from litellm.llms.vertex_ai.common_utils import process_items - # Test empty items + # Test empty items — should preserve "any type" semantics (not coerce to object) schema = {"type": "array", "items": {}} process_items(schema) - assert schema["items"] == {"type": "object"} + assert schema["items"] == {} - # Test nested items + # Test nested items — should preserve "any type" semantics schema = {"type": "array", "items": {"type": "array", "items": {}}} process_items(schema) - assert schema["items"]["items"] == {"type": "object"} + assert schema["items"]["items"] == {} - # Test items in properties + # Test items in properties — should preserve "any type" semantics schema = { "type": "object", "properties": {"nested": {"type": "array", "items": {}}}, } process_items(schema) - assert schema["properties"]["nested"]["items"] == {"type": "object"} + assert schema["properties"]["nested"]["items"] == {} + + # Test items with actual type — should not be modified + schema = {"type": "array", "items": {"type": "string"}} + process_items(schema) + assert schema["items"] == {"type": "string"} def test_vertex_ai_complex_response_schema(): @@ -1402,3 +1407,89 @@ def test_add_object_type_does_not_add_type_when_anyof_present(): # Verify type was not added (anyOf handles the type) assert "type" not in input_schema, "type should not be added when anyOf is present" + + +def test_is_any_type_schema(): + """Test _is_any_type_schema correctly identifies unconstrained schemas.""" + from litellm.llms.vertex_ai.common_utils import _is_any_type_schema + + # Empty schema = any type + assert _is_any_type_schema({}) is True + + # Only metadata keys = any type + assert _is_any_type_schema({"description": "Any value"}) is True + assert _is_any_type_schema({"title": "MyField"}) is True + assert _is_any_type_schema({"title": "X", "description": "Y", "default": 0}) is True + + # Has type-constraining keys = NOT any type + assert _is_any_type_schema({"type": "object"}) is False + assert _is_any_type_schema({"type": "string"}) is False + assert _is_any_type_schema({"properties": {"a": {}}}) is False + assert _is_any_type_schema({"items": {"type": "string"}}) is False + assert _is_any_type_schema({"anyOf": [{"type": "string"}]}) is False + assert _is_any_type_schema({"$schema": "https://json-schema.org/draft/2020-12/schema"}) is False + assert _is_any_type_schema({"enum": ["a", "b"]}) is False + + +def test_add_object_type_preserves_any_type_schema(): + """Test add_object_type does NOT add type:object to empty schemas (any type).""" + from litellm.llms.vertex_ai.common_utils import add_object_type + + # Empty schema should be preserved (any type) + schema = {} + add_object_type(schema) + assert "type" not in schema, "Empty schema (any type) should not get type: object" + + # Schema with only description should be preserved + schema = {"description": "Any JSON value"} + add_object_type(schema) + assert "type" not in schema + + # Schema with $schema key should still get type: object (tool with no args) + schema = {"$schema": "https://json-schema.org/draft/2020-12/schema"} + add_object_type(schema) + assert schema["type"] == "object" + + +def test_convert_anyof_preserves_any_type_members(): + """Test convert_anyof_null_to_nullable does NOT coerce empty anyOf members to object.""" + from litellm.llms.vertex_ai.common_utils import convert_anyof_null_to_nullable + + # anyOf with empty schema and null — empty should be preserved + schema = { + "anyOf": [ + {}, + {"type": "null"}, + ] + } + convert_anyof_null_to_nullable(schema) + # null should be removed, empty schema should be preserved (not coerced to object) + assert len(schema["anyOf"]) == 1 + assert "type" not in schema["anyOf"][0] or schema["anyOf"][0].get("type") != "object" + assert schema["anyOf"][0].get("nullable") is True + + +def test_build_vertex_schema_jsonvalue(): + """ + End-to-end: Pydantic JsonValue generates {} in $defs. + _build_vertex_schema should preserve any-type semantics. + Regression test for https://github.com/BerriAI/litellm/issues/22391 + """ + from litellm.llms.vertex_ai.common_utils import _build_vertex_schema + + # Simulates what Pydantic generates for a model with JsonValue field + schema = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "value": {}, # after $ref resolution, this is what JsonValue becomes + }, + "required": ["name", "value"], + } + result = _build_vertex_schema(schema) + + # The "value" field should NOT have been coerced to type: object + value_schema = result["properties"]["value"] + assert value_schema.get("type") != "object", ( + "JsonValue schema {} should not be coerced to {type: object}" + ) From 0e024221a81875bb656d050e9f58b33bc39b7d71 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 2 Mar 2026 17:05:37 -0300 Subject: [PATCH 028/147] refactor: remove no-op pass branch in process_items MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Greptile review nit — the if + pass was a no-op since the mutation was removed. Deleting the branch entirely. --- litellm/llms/vertex_ai/common_utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 8f8baad94c8..c2f4257cd45 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -615,8 +615,6 @@ def process_items(schema, depth=0): f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting." ) if isinstance(schema, dict): - if "items" in schema and isinstance(schema["items"], dict) and _is_any_type_schema(schema["items"]): - pass # preserve "any type" semantics — don't coerce to object for key, value in schema.items(): if isinstance(value, dict): process_items(value, depth + 1) From 7d23106fcfaf304045a3ce976b50d0783f807215 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 2 Mar 2026 19:15:37 -0300 Subject: [PATCH 029/147] fix(helicone): correct provider URL for Vertex AI Gemini models Reorder elif branches so is_vertex_ai is checked before "gemini" in model. Previously, Vertex AI Gemini models (e.g. vertex_ai/gemini-2.5-flash) matched the "gemini" substring check first and were logged with the Google AI Studio URL instead of the Vertex AI URL. --- litellm/integrations/helicone.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py index c77a1b2564a..51e6699c5f4 100644 --- a/litellm/integrations/helicone.py +++ b/litellm/integrations/helicone.py @@ -167,12 +167,12 @@ class HeliconeLogger: if "claude" in model and not is_vertex_ai: url = f"{self.api_base}/anthropic/v1/log" provider_url = "https://api.anthropic.com/v1/messages" - elif "gemini" in model: - url = f"{self.api_base}/custom/v1/log" - provider_url = "https://generativelanguage.googleapis.com/v1beta" elif is_vertex_ai: url = f"{self.api_base}/custom/v1/log" provider_url = "https://aiplatform.googleapis.com/v1" + elif "gemini" in model: + url = f"{self.api_base}/custom/v1/log" + provider_url = "https://generativelanguage.googleapis.com/v1beta" headers = { "Authorization": f"Bearer {self.key}", "Content-Type": "application/json", From 81ddf084943190f2725c758b81c4ac285ac9f569 Mon Sep 17 00:00:00 2001 From: liweiguang Date: Tue, 3 Mar 2026 12:12:43 +0800 Subject: [PATCH 030/147] fix: add missing `supports_function_calling` for deepinfra models All 55 deepinfra models that had `supports_tool_choice: true` were missing the `supports_function_calling` flag, causing `litellm.supports_function_calling()` to incorrectly return False. Fixes #22619 Co-Authored-By: Claude Opus 4.6 --- model_prices_and_context_window.json | 166 ++++++++++++++++++--------- tests/test_litellm/test_utils.py | 13 +++ 2 files changed, 124 insertions(+), 55 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4934f11d456..a631f0b5b8e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10844,7 +10844,8 @@ "output_cost_per_token": 9e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": { "max_tokens": 131072, @@ -10854,7 +10855,8 @@ "output_cost_per_token": 1e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": { "max_tokens": 131072, @@ -10874,7 +10876,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen2.5-72B-Instruct": { "max_tokens": 32768, @@ -10884,7 +10887,8 @@ "output_cost_per_token": 3.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen2.5-7B-Instruct": { "max_tokens": 32768, @@ -10905,7 +10909,8 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-14B": { "max_tokens": 40960, @@ -10915,7 +10920,8 @@ "output_cost_per_token": 2.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B": { "max_tokens": 40960, @@ -10925,7 +10931,8 @@ "output_cost_per_token": 5.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": { "max_tokens": 262144, @@ -10935,7 +10942,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": { "max_tokens": 262144, @@ -10945,7 +10953,8 @@ "output_cost_per_token": 2.9e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-30B-A3B": { "max_tokens": 40960, @@ -10955,7 +10964,8 @@ "output_cost_per_token": 2.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-32B": { "max_tokens": 40960, @@ -10965,7 +10975,8 @@ "output_cost_per_token": 2.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": { "max_tokens": 262144, @@ -10975,7 +10986,8 @@ "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": { "max_tokens": 262144, @@ -10985,7 +10997,8 @@ "output_cost_per_token": 1.2e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Next-80B-A3B-Instruct": { "max_tokens": 262144, @@ -10995,7 +11008,8 @@ "output_cost_per_token": 1.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Next-80B-A3B-Thinking": { "max_tokens": 262144, @@ -11005,7 +11019,8 @@ "output_cost_per_token": 1.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": { "max_tokens": 8192, @@ -11056,7 +11071,8 @@ "cache_read_input_token_cost": 3.3e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/anthropic/claude-4-opus": { "max_tokens": 200000, @@ -11066,7 +11082,8 @@ "output_cost_per_token": 8.25e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/anthropic/claude-4-sonnet": { "max_tokens": 200000, @@ -11076,7 +11093,8 @@ "output_cost_per_token": 1.65e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1": { "max_tokens": 163840, @@ -11086,7 +11104,8 @@ "output_cost_per_token": 2.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528": { "max_tokens": 163840, @@ -11097,7 +11116,8 @@ "cache_read_input_token_cost": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": { "max_tokens": 32768, @@ -11107,7 +11127,8 @@ "output_cost_per_token": 3e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { "max_tokens": 131072, @@ -11127,7 +11148,8 @@ "output_cost_per_token": 2.7e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": { "max_tokens": 40960, @@ -11137,7 +11159,8 @@ "output_cost_per_token": 3e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3": { "max_tokens": 163840, @@ -11147,7 +11170,8 @@ "output_cost_per_token": 8.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3-0324": { "max_tokens": 163840, @@ -11157,7 +11181,8 @@ "output_cost_per_token": 8.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1": { "max_tokens": 163840, @@ -11169,7 +11194,8 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_reasoning": true + "supports_reasoning": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1-Terminus": { "max_tokens": 163840, @@ -11180,7 +11206,8 @@ "cache_read_input_token_cost": 2.16e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.0-flash-001": { "deprecation_date": "2026-06-01", @@ -11191,7 +11218,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.5-flash": { "max_tokens": 1000000, @@ -11201,7 +11229,8 @@ "output_cost_per_token": 2.5e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -11211,7 +11240,8 @@ "output_cost_per_token": 1e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -11221,7 +11251,8 @@ "output_cost_per_token": 1e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-27b-it": { "max_tokens": 131072, @@ -11231,7 +11262,8 @@ "output_cost_per_token": 1.6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-4b-it": { "max_tokens": 131072, @@ -11241,7 +11273,8 @@ "output_cost_per_token": 8e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": { "max_tokens": 131072, @@ -11261,7 +11294,8 @@ "output_cost_per_token": 2e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct": { "max_tokens": 131072, @@ -11271,7 +11305,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": { "max_tokens": 131072, @@ -11281,6 +11316,7 @@ "output_cost_per_token": 3.9e-07, "litellm_provider": "deepinfra", "mode": "chat", + "supports_function_calling": true, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { @@ -11291,7 +11327,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": { "max_tokens": 327680, @@ -11301,7 +11338,8 @@ "output_cost_per_token": 3e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-Guard-3-8B": { "max_tokens": 131072, @@ -11331,7 +11369,8 @@ "output_cost_per_token": 6e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": { "max_tokens": 131072, @@ -11341,7 +11380,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { "max_tokens": 131072, @@ -11351,7 +11391,8 @@ "output_cost_per_token": 2.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": { "max_tokens": 131072, @@ -11361,7 +11402,8 @@ "output_cost_per_token": 5e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { "max_tokens": 131072, @@ -11371,7 +11413,8 @@ "output_cost_per_token": 3e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/microsoft/WizardLM-2-8x22B": { "max_tokens": 65536, @@ -11391,7 +11434,8 @@ "output_cost_per_token": 1.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": { "max_tokens": 131072, @@ -11401,7 +11445,8 @@ "output_cost_per_token": 4e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": { "max_tokens": 32768, @@ -11411,7 +11456,8 @@ "output_cost_per_token": 8e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": { "max_tokens": 128000, @@ -11421,7 +11467,8 @@ "output_cost_per_token": 2e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { "max_tokens": 32768, @@ -11431,7 +11478,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/moonshotai/Kimi-K2-Instruct": { "max_tokens": 131072, @@ -11441,7 +11489,8 @@ "output_cost_per_token": 2e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/moonshotai/Kimi-K2-Instruct-0905": { "max_tokens": 262144, @@ -11452,7 +11501,8 @@ "cache_read_input_token_cost": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": { "max_tokens": 131072, @@ -11462,7 +11512,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/Llama-3.3-Nemotron-Super-49B-v1.5": { "max_tokens": 131072, @@ -11472,7 +11523,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { "max_tokens": 131072, @@ -11482,7 +11534,8 @@ "output_cost_per_token": 1.6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/openai/gpt-oss-120b": { "max_tokens": 131072, @@ -11492,7 +11545,8 @@ "output_cost_per_token": 4.5e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/openai/gpt-oss-20b": { "max_tokens": 131072, @@ -11502,7 +11556,8 @@ "output_cost_per_token": 1.5e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/zai-org/GLM-4.5": { "max_tokens": 131072, @@ -11512,7 +11567,8 @@ "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepseek/deepseek-chat": { "cache_creation_input_token_cost": 0.0, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 7f0b3b5b501..c0777559584 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -96,6 +96,19 @@ def test_supports_function_calling_github_anthropic_alias(): ) +def test_supports_function_calling_deepinfra_llama(): + """Test that deepinfra Llama models correctly report function calling support. + + Regression test for https://github.com/BerriAI/litellm/issues/22619 + """ + assert ( + litellm.utils.supports_function_calling( + model="deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo" + ) + is True + ) + + def test_supports_function_calling_unknown_github_alias_returns_false(): assert ( litellm.utils.supports_function_calling( From 18216ac07c68dee1d9fd8d1d6f4788b62b9394ec Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 10:48:00 +0530 Subject: [PATCH 031/147] Fix: Azure ai finetuning api --- litellm/fine_tuning/main.py | 55 +++ ...odel_prices_and_context_window_backup.json | 353 +++++++++++++++--- litellm/types/llms/openai.py | 13 +- tests/batches_tests/test_fine_tuning_api.py | 58 +++ 4 files changed, 423 insertions(+), 56 deletions(-) diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index f5b8b097026..db77fa32919 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -34,6 +34,44 @@ vertex_fine_tuning_apis_instance = VertexFineTuningAPI() ################################################# +def _prepare_azure_extra_body( + extra_body: Optional[Dict[str, Any]], + kwargs: Dict[str, Any], + azure_specific_hyperparams: Dict[str, Any], +) -> Dict[str, Any]: + """ + Prepare extra_body for Azure fine-tuning API by combining Azure-specific parameters. + + Azure fine-tuning API accepts additional parameters beyond the standard OpenAI spec: + - trainingType: Type of training (e.g., 1 for supervised fine-tuning) + - prompt_loss_weight: Weight for prompt loss in training + + These parameters must be passed in the extra_body field when calling the Azure OpenAI SDK. + + Args: + extra_body: Optional existing extra_body dict + kwargs: Request kwargs that may contain Azure-specific parameters + azure_specific_hyperparams: Dict of Azure-specific hyperparameters already extracted + + Returns: + Dict containing all Azure-specific parameters to be passed in extra_body + """ + if extra_body is None: + extra_body = {} + + # Azure-specific root-level parameters + azure_specific_params = ["trainingType"] + for param in azure_specific_params: + if param in kwargs: + extra_body[param] = kwargs[param] + + # Add Azure-specific hyperparameters + if azure_specific_hyperparams: + extra_body.update(azure_specific_hyperparams) + + return extra_body + + @client async def acreate_fine_tuning_job( model: str, @@ -114,6 +152,15 @@ def create_fine_tuning_job( # handle hyperparameters hyperparameters = hyperparameters or {} # original hyperparameters + + # For Azure, extract Azure-specific hyperparameters before creating OpenAI-spec hyperparameters + azure_specific_hyperparams = {} + if custom_llm_provider == "azure": + azure_hyperparameter_keys = ["prompt_loss_weight"] + for key in azure_hyperparameter_keys: + if key in hyperparameters: + azure_specific_hyperparams[key] = hyperparameters.pop(key) + _oai_hyperparameters: Hyperparameters = Hyperparameters( **hyperparameters ) # Typed Hyperparameters for OpenAI Spec @@ -207,6 +254,10 @@ def create_fine_tuning_job( extra_body.pop("azure_ad_token", None) else: get_secret_str("AZURE_AD_TOKEN") # type: ignore + + # Prepare Azure-specific parameters for extra_body + extra_body = _prepare_azure_extra_body(extra_body, kwargs, azure_specific_hyperparams) + create_fine_tuning_job_data = FineTuningJobCreate( model=model, training_file=training_file, @@ -220,6 +271,10 @@ def create_fine_tuning_job( create_fine_tuning_job_data_dict = create_fine_tuning_job_data.model_dump( exclude_none=True ) + + # Add extra_body if it has Azure-specific parameters + if extra_body: + create_fine_tuning_job_data_dict["extra_body"] = extra_body response = azure_fine_tuning_apis_instance.create_fine_tuning_job( api_base=api_base, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d4c5b476af6..4934f11d456 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -846,7 +846,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "anthropic.claude-3-opus-20240229-v1:0": { "input_cost_per_token": 1.5e-05, @@ -859,7 +861,9 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost": 1.875e-05 }, "anthropic.claude-3-sonnet-20240229-v1:0": { "input_cost_per_token": 3e-06, @@ -873,7 +877,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "anthropic.claude-instant-v1": { "input_cost_per_token": 8e-07, @@ -1512,7 +1518,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "apac.anthropic.claude-3-5-sonnet-20241022-v2:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -1545,7 +1553,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "apac.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, @@ -1581,7 +1591,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "apac.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -6925,7 +6937,9 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "bedrock/sa-east-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 4.45e-06, @@ -7344,7 +7358,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3.6e-07, + "cache_creation_input_token_cost": 4.5e-06 }, "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0": { "input_cost_per_token": 3e-07, @@ -7358,7 +7374,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-08, + "cache_creation_input_token_cost": 3.75e-07 }, "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0": { "input_cost_per_token": 3.3e-06, @@ -7376,7 +7394,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost": 4.125e-06 }, "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -7489,7 +7509,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3.6e-07, + "cache_creation_input_token_cost": 4.5e-06 }, "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0": { "input_cost_per_token": 3e-07, @@ -7503,7 +7525,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-08, + "cache_creation_input_token_cost": 3.75e-07 }, "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0": { "input_cost_per_token": 3.3e-06, @@ -7521,7 +7545,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost": 4.125e-06 }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -9753,6 +9779,74 @@ } ] }, + "dashscope/qwen3-vl-plus": { + "litellm_provider": "dashscope", + "max_input_tokens": 260096, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 4.8e-06, + "range": [ + 128000.0, + 256000.0 + ] + } + ] + }, + "dashscope/qwen3.5-plus": { + "litellm_provider": "dashscope", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 3e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, "dashscope/qwq-plus": { "input_cost_per_token": 8e-07, "litellm_provider": "dashscope", @@ -11089,7 +11183,7 @@ "supports_tool_choice": true }, "deepinfra/google/gemini-2.0-flash-001": { - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -11950,7 +12044,9 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, @@ -11987,7 +12083,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "eu.anthropic.claude-3-5-sonnet-20241022-v2:0": { "input_cost_per_token": 3e-06, @@ -12004,7 +12102,9 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "eu.anthropic.claude-3-7-sonnet-20250219-v1:0": { "input_cost_per_token": 3e-06, @@ -12022,7 +12122,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "eu.anthropic.claude-3-haiku-20240307-v1:0": { "input_cost_per_token": 2.5e-07, @@ -12036,7 +12138,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "eu.anthropic.claude-3-opus-20240229-v1:0": { "input_cost_per_token": 1.5e-05, @@ -12049,7 +12153,9 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost": 1.875e-05 }, "eu.anthropic.claude-3-sonnet-20240229-v1:0": { "input_cost_per_token": 3e-06, @@ -12063,7 +12169,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "eu.anthropic.claude-opus-4-1-20250805-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -13590,7 +13698,7 @@ }, "gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -13630,7 +13738,7 @@ }, "gemini-2.0-flash-001": { "cache_read_input_token_cost": 3.75e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-language-models", @@ -13716,7 +13824,7 @@ }, "gemini-2.0-flash-lite": { "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, "litellm_provider": "vertex_ai-language-models", @@ -13752,7 +13860,7 @@ }, "gemini-2.0-flash-lite-001": { "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, "litellm_provider": "vertex_ai-language-models", @@ -14669,6 +14777,7 @@ "supports_web_search": true }, "gemini-3-pro-preview": { + "deprecation_date": "2026-03-26", "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -15805,7 +15914,7 @@ }, "gemini/gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -15846,7 +15955,7 @@ }, "gemini/gemini-2.0-flash-001": { "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -15934,7 +16043,7 @@ }, "gemini/gemini-2.0-flash-lite": { "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, "litellm_provider": "gemini", @@ -15970,7 +16079,7 @@ "tpm": 4000000 }, "gemini/gemini-2.0-flash-lite-preview-02-05": { - "deprecation_date": "2025-12-02", + "deprecation_date": "2025-12-09", "cache_read_input_token_cost": 1.875e-08, "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, @@ -16925,6 +17034,7 @@ "tpm": 800000 }, "gemini/gemini-3-pro-preview": { + "deprecation_date": "2026-03-09", "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 2e-06, @@ -23112,6 +23222,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/magistral-medium-1-2-2509": { + "input_cost_per_token": 2e-06, + "litellm_provider": "mistral", + "max_input_tokens": 40000, + "max_output_tokens": 40000, + "max_tokens": 40000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://mistral.ai/news/magistral", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.001, @@ -23177,6 +23302,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/magistral-small-1-2-2509": { + "input_cost_per_token": 5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 40000, + "max_output_tokens": 40000, + "max_tokens": 40000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://mistral.ai/pricing#api-pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/mistral-embed": { "input_cost_per_token": 1e-07, "litellm_provider": "mistral", @@ -23238,24 +23378,41 @@ "supports_tool_choice": true }, "mistral/mistral-large-latest": { - "input_cost_per_token": 2e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "mistral", - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 1.5e-06, + "source": "https://docs.mistral.ai/models/mistral-large-3-25-12", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-large-3": { "input_cost_per_token": 5e-07, "litellm_provider": "mistral", - "max_input_tokens": 256000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.mistral.ai/models/mistral-large-3-25-12", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-large-2512": { + "input_cost_per_token": 5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://docs.mistral.ai/models/mistral-large-3-25-12", @@ -23306,14 +23463,30 @@ "input_cost_per_token": 4e-07, "litellm_provider": "mistral", "max_input_tokens": 131072, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-3-1-2508": { + "input_cost_per_token": 4e-07, + "litellm_provider": "mistral", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/mistral-medium-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-small": { "input_cost_per_token": 1e-07, @@ -23329,17 +23502,79 @@ "supports_tool_choice": true }, "mistral/mistral-small-latest": { - "input_cost_per_token": 1e-07, + "input_cost_per_token": 6e-08, "litellm_provider": "mistral", - "max_input_tokens": 32000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-07, + "output_cost_per_token": 1.8e-07, + "source": "https://mistral.ai/pricing", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-small-3-2-2506": { + "input_cost_per_token": 6e-08, + "litellm_provider": "mistral", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.8e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/ministral-3-3b-2512": { + "input_cost_per_token": 1e-07, + "litellm_provider": "mistral", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/ministral-3-8b-2512": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/ministral-3-14b-2512": { + "input_cost_per_token": 2e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-tiny": { "input_cost_per_token": 2.5e-07, @@ -25657,7 +25892,7 @@ "supports_tool_choice": true }, "openrouter/google/gemini-2.0-flash-001": { - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", @@ -29554,7 +29789,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "us.anthropic.claude-3-5-sonnet-20241022-v2:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -29607,7 +29844,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "us.anthropic.claude-3-opus-20240229-v1:0": { "input_cost_per_token": 1.5e-05, @@ -29620,7 +29859,9 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost": 1.875e-05 }, "us.anthropic.claude-3-sonnet-20240229-v1:0": { "input_cost_per_token": 3e-06, @@ -29634,7 +29875,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "us.anthropic.claude-opus-4-1-20250805-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -30527,7 +30770,7 @@ "supports_tool_choice": true }, "vercel_ai_gateway/google/gemini-2.0-flash": { - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_token": 1.5e-07, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, @@ -30541,7 +30784,7 @@ "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.0-flash-lite": { - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_token": 7.5e-08, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, @@ -37898,7 +38141,7 @@ }, "gemini/gemini-2.0-flash-lite-001": { "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, "litellm_provider": "gemini", diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index f82f6a02f22..d06d879dad1 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -71,7 +71,14 @@ from openai.types.responses.response_create_params import ( ToolParam, ) from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from pydantic import BaseModel, ConfigDict, Discriminator, PrivateAttr, field_serializer, field_validator +from pydantic import ( + BaseModel, + ConfigDict, + Discriminator, + PrivateAttr, + field_serializer, + field_validator, +) from typing_extensions import Annotated, Dict, Required, TypedDict, override from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject @@ -964,6 +971,10 @@ class Hyperparameters(BaseModel): n_epochs: Optional[Union[str, int]] = ( None # "The number of epochs to train the model for" ) + + model_config = { + "extra": "allow" + } class FineTuningJobCreate(BaseModel): diff --git a/tests/batches_tests/test_fine_tuning_api.py b/tests/batches_tests/test_fine_tuning_api.py index c6a731ea54f..7e238173480 100644 --- a/tests/batches_tests/test_fine_tuning_api.py +++ b/tests/batches_tests/test_fine_tuning_api.py @@ -596,3 +596,61 @@ async def test_mock_openai_retrieve_fine_tune_job(): # Verify the request mock_retrieve.assert_called_once_with(fine_tuning_job_id="ft-123") + + +@pytest.mark.asyncio +async def test_mock_azure_create_fine_tune_job_with_azure_specific_params(): + """Test that Azure-specific parameters are passed through extra_body""" + from openai import AsyncAzureOpenAI + from openai.types.fine_tuning.fine_tuning_job import FineTuningJob + from openai.types.fine_tuning.fine_tuning_job import Hyperparameters as OAIHyperparameters + + mock_response = FineTuningJob( + id="ft-azure-123", + model="gpt-4.1-mini-2025-04-14", + created_at=1677610602, + status="validating_files", + fine_tuned_model=None, + object="fine_tuning.job", + hyperparameters=OAIHyperparameters(n_epochs=3), + organization_id="org-123", + seed=42, + training_file="file-123", + result_files=[], + ) + + with patch("litellm.llms.azure.fine_tuning.handler.AzureOpenAIFineTuningAPI.create_fine_tuning_job") as mock_create: + mock_create.return_value = mock_response + + response = await litellm.acreate_fine_tuning_job( + model="gpt-4.1-mini-2025-04-14", + training_file="file-123", + custom_llm_provider="azure", + api_base="https://test.openai.azure.com", + api_key="test-key", + api_version="2025-04-01-preview", + trainingType=1, + hyperparameters={ + "n_epochs": 3, + "prompt_loss_weight": 0.1 + }, + ) + + # Verify the request + mock_create.assert_called_once() + request_params = mock_create.call_args.kwargs + + # Check that create_fine_tuning_job_data contains the correct structure + create_data = request_params["create_fine_tuning_job_data"] + assert create_data["model"] == "gpt-4.1-mini-2025-04-14" + assert create_data["training_file"] == "file-123" + assert create_data["hyperparameters"] == {"n_epochs": 3} + + # Azure-specific parameters should be in extra_body + assert "extra_body" in create_data + assert create_data["extra_body"]["trainingType"] == 1 + assert create_data["extra_body"]["prompt_loss_weight"] == 0.1 + + # Verify the response + assert response.id == "ft-azure-123" + assert response.model == "gpt-4.1-mini-2025-04-14" From d07689d2d70804959dbd2e2e39cc797e3a96e37e Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 3 Mar 2026 11:59:58 +0530 Subject: [PATCH 032/147] =?UTF-8?q?bump:=20version=201.82.0=20=E2=86=92=20?= =?UTF-8?q?1.82.1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 577e51a0d22..6f9add2e4cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.82.0" +version = "1.82.1" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -183,7 +183,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.82.0" +version = "1.82.1" version_files = [ "pyproject.toml:^version" ] From f8034f15ada8daf1bb0ed08ead330e3f7e67eee8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 14:51:12 +0530 Subject: [PATCH 033/147] Remove defualt hardcoded thinking levels for gemini 3 family --- .../vertex_and_google_ai_studio_gemini.py | 17 ----------------- 1 file changed, 17 deletions(-) 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 dee826e5783..0905f22362e 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 @@ -1136,23 +1136,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if VertexGeminiConfig._is_gemini_3_or_newer(model): if "temperature" not in optional_params: optional_params["temperature"] = 1.0 - # Only add thinkingLevel if model supports it (exclude image models) - if "image" not in model.lower(): - thinking_config = optional_params.get("thinkingConfig", {}) - if ( - "thinkingLevel" not in thinking_config - and "thinkingBudget" not in thinking_config - ): - # For gemini-3-flash-preview, default to "minimal" to match Gemini 2.5 Flash behavior - # For other Gemini 3 models, default to "low" - is_gemini3flash = ( - "gemini-3-flash-preview" in model.lower() - or "gemini-3-flash" in model.lower() - ) - thinking_config["thinkingLevel"] = ( - "minimal" if is_gemini3flash else "low" - ) - optional_params["thinkingConfig"] = thinking_config return optional_params From 213423cb45308e47538b207dc043c1149e75cda7 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 15:05:20 +0530 Subject: [PATCH 034/147] Fix test case --- .../test_vertex_and_google_ai_studio_gemini.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) 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 196bb00f40d..8beb19bf1ac 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 @@ -2130,7 +2130,7 @@ def test_reasoning_effort_dict_format_gemini_3(): assert result["thinkingConfig"]["thinkingLevel"] == "high" assert result["thinkingConfig"]["includeThoughts"] is True - # Test dict format without effort key - should fall back to Gemini 3 default (low) + # Test dict format without effort key - no thinkingConfig should be set optional_params = {} non_default_params = {"reasoning_effort": {"summary": "auto"}} result = v.map_openai_params( @@ -2139,8 +2139,8 @@ def test_reasoning_effort_dict_format_gemini_3(): model=model, drop_params=False, ) - # Gemini 3 defaults to thinkingLevel="low" when no explicit effort is set - assert result["thinkingConfig"]["thinkingLevel"] == "low" + # No effort key in dict → no thinkingConfig set + assert "thinkingConfig" not in result def test_temperature_default_for_gemini_3(): @@ -2453,8 +2453,8 @@ def test_gemini_3_image_models_no_thinking_config(): def test_gemini_3_text_models_get_thinking_config(): """ - Test that Gemini 3 text models DO receive automatic thinkingConfig. - This ensures we didn't break the existing behavior for non-image models. + Test that Gemini 3 text models do NOT receive automatic thinkingConfig + when no reasoning_effort or thinking param is provided. """ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -2462,7 +2462,7 @@ def test_gemini_3_text_models_get_thinking_config(): v = VertexGeminiConfig() - # Test gemini-3-pro-preview (text model, should get thinking) + # Test gemini-3-pro-preview (text model, no explicit thinking params) model = "gemini-3-pro-preview" optional_params = {} non_default_params = {} @@ -2474,9 +2474,8 @@ def test_gemini_3_text_models_get_thinking_config(): drop_params=False, ) - # Should have thinkingConfig automatically added - assert "thinkingConfig" in result - assert result["thinkingConfig"]["thinkingLevel"] == "low" + # Should NOT have thinkingConfig automatically added when user provides no reasoning_effort + assert "thinkingConfig" not in result assert result["temperature"] == 1.0 From 851be587751fef0dca958394c9bd77746de1cae6 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 15:07:47 +0530 Subject: [PATCH 035/147] Add day 0 support of gemini-3.1-flash-lite-preview --- docs/my-website/docs/providers/gemini.md | 1 + docs/my-website/docs/providers/vertex.md | 1 + ...odel_prices_and_context_window_backup.json | 498 ++++++++++++++++-- model_prices_and_context_window.json | 145 +++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 48 +- ...test_vertex_and_google_ai_studio_gemini.py | 11 +- 6 files changed, 642 insertions(+), 62 deletions(-) diff --git a/docs/my-website/docs/providers/gemini.md b/docs/my-website/docs/providers/gemini.md index 6de2263916c..f97f025c19b 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/docs/my-website/docs/providers/gemini.md @@ -2041,6 +2041,7 @@ response = litellm.completion( | gemini-2.0-flash-lite-preview-02-05 | `completion(model='gemini/gemini-2.0-flash-lite-preview-02-05', messages)` | `os.environ['GEMINI_API_KEY']` | | gemini-2.5-flash-preview-09-2025 | `completion(model='gemini/gemini-2.5-flash-preview-09-2025', messages)` | `os.environ['GEMINI_API_KEY']` | | gemini-2.5-flash-lite-preview-09-2025 | `completion(model='gemini/gemini-2.5-flash-lite-preview-09-2025', messages)` | `os.environ['GEMINI_API_KEY']` | +| gemini-3.1-flash-lite-preview | `completion(model='gemini/gemini-3.1-flash-lite-preview', messages)` | `os.environ['GEMINI_API_KEY']` | | gemini-flash-latest | `completion(model='gemini/gemini-flash-latest', messages)` | `os.environ['GEMINI_API_KEY']` | | gemini-flash-lite-latest | `completion(model='gemini/gemini-flash-lite-latest', messages)` | `os.environ['GEMINI_API_KEY']` | diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 63e4dceec00..94619082e88 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -1685,6 +1685,7 @@ litellm.vertex_location = "us-central1 # Your Location | gemini-2.5-pro | `completion('gemini-2.5-pro', messages)`, `completion('vertex_ai/gemini-2.5-pro', messages)` | | gemini-2.5-flash-preview-09-2025 | `completion('gemini-2.5-flash-preview-09-2025', messages)`, `completion('vertex_ai/gemini-2.5-flash-preview-09-2025', messages)` | | gemini-2.5-flash-lite-preview-09-2025 | `completion('gemini-2.5-flash-lite-preview-09-2025', messages)`, `completion('vertex_ai/gemini-2.5-flash-lite-preview-09-2025', messages)` | +| gemini-3.1-flash-lite-preview | `completion('gemini-3.1-flash-lite-preview', messages)`, `completion('vertex_ai/gemini-3.1-flash-lite-preview', messages)` | ## Private Service Connect (PSC) Endpoints diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d4c5b476af6..42b4f0f7762 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -846,7 +846,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "anthropic.claude-3-opus-20240229-v1:0": { "input_cost_per_token": 1.5e-05, @@ -859,7 +861,9 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost": 1.875e-05 }, "anthropic.claude-3-sonnet-20240229-v1:0": { "input_cost_per_token": 3e-06, @@ -873,7 +877,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "anthropic.claude-instant-v1": { "input_cost_per_token": 8e-07, @@ -1512,7 +1518,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "apac.anthropic.claude-3-5-sonnet-20241022-v2:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -1545,7 +1553,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "apac.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, @@ -1581,7 +1591,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "apac.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -6925,7 +6937,9 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "bedrock/sa-east-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 4.45e-06, @@ -7344,7 +7358,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3.6e-07, + "cache_creation_input_token_cost": 4.5e-06 }, "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0": { "input_cost_per_token": 3e-07, @@ -7358,7 +7374,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-08, + "cache_creation_input_token_cost": 3.75e-07 }, "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0": { "input_cost_per_token": 3.3e-06, @@ -7376,7 +7394,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost": 4.125e-06 }, "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -7489,7 +7509,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3.6e-07, + "cache_creation_input_token_cost": 4.5e-06 }, "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0": { "input_cost_per_token": 3e-07, @@ -7503,7 +7525,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-08, + "cache_creation_input_token_cost": 3.75e-07 }, "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0": { "input_cost_per_token": 3.3e-06, @@ -7521,7 +7545,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost": 4.125e-06 }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { "input_cost_per_token": 2.65e-06, @@ -9753,6 +9779,74 @@ } ] }, + "dashscope/qwen3-vl-plus": { + "litellm_provider": "dashscope", + "max_input_tokens": 260096, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 4.8e-06, + "range": [ + 128000.0, + 256000.0 + ] + } + ] + }, + "dashscope/qwen3.5-plus": { + "litellm_provider": "dashscope", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 3e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, "dashscope/qwq-plus": { "input_cost_per_token": 8e-07, "litellm_provider": "dashscope", @@ -11089,7 +11183,7 @@ "supports_tool_choice": true }, "deepinfra/google/gemini-2.0-flash-001": { - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -11950,7 +12044,9 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "eu.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, @@ -11987,7 +12083,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "eu.anthropic.claude-3-5-sonnet-20241022-v2:0": { "input_cost_per_token": 3e-06, @@ -12004,7 +12102,9 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "eu.anthropic.claude-3-7-sonnet-20250219-v1:0": { "input_cost_per_token": 3e-06, @@ -12022,7 +12122,9 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "eu.anthropic.claude-3-haiku-20240307-v1:0": { "input_cost_per_token": 2.5e-07, @@ -12036,7 +12138,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "eu.anthropic.claude-3-opus-20240229-v1:0": { "input_cost_per_token": 1.5e-05, @@ -12049,7 +12153,9 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost": 1.875e-05 }, "eu.anthropic.claude-3-sonnet-20240229-v1:0": { "input_cost_per_token": 3e-06, @@ -12063,7 +12169,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "eu.anthropic.claude-opus-4-1-20250805-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -13590,7 +13698,7 @@ }, "gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -13630,7 +13738,7 @@ }, "gemini-2.0-flash-001": { "cache_read_input_token_cost": 3.75e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-language-models", @@ -13716,7 +13824,7 @@ }, "gemini-2.0-flash-lite": { "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, "litellm_provider": "vertex_ai-language-models", @@ -13752,7 +13860,7 @@ }, "gemini-2.0-flash-lite-001": { "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, "litellm_provider": "vertex_ai-language-models", @@ -14226,6 +14334,53 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-09, + "input_cost_per_audio_token": 2.5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true + }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -14669,6 +14824,7 @@ "supports_web_search": true }, "gemini-3-pro-preview": { + "deprecation_date": "2026-03-26", "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -15805,7 +15961,7 @@ }, "gemini/gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -15846,7 +16002,7 @@ }, "gemini/gemini-2.0-flash-001": { "cache_read_input_token_cost": 2.5e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -15934,7 +16090,7 @@ }, "gemini/gemini-2.0-flash-lite": { "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, "litellm_provider": "gemini", @@ -15970,7 +16126,7 @@ "tpm": 4000000 }, "gemini/gemini-2.0-flash-lite-preview-02-05": { - "deprecation_date": "2025-12-02", + "deprecation_date": "2025-12-09", "cache_read_input_token_cost": 1.875e-08, "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, @@ -16925,6 +17081,7 @@ "tpm": 800000 }, "gemini/gemini-3-pro-preview": { + "deprecation_date": "2026-03-09", "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 2e-06, @@ -16980,6 +17137,56 @@ "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "supports_service_tier": true }, + "gemini/gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-09, + "input_cost_per_audio_token": 2.5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 250000 + }, "gemini/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, "input_cost_per_audio_token": 1e-06, @@ -23112,6 +23319,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/magistral-medium-1-2-2509": { + "input_cost_per_token": 2e-06, + "litellm_provider": "mistral", + "max_input_tokens": 40000, + "max_output_tokens": 40000, + "max_tokens": 40000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://mistral.ai/news/magistral", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.001, @@ -23177,6 +23399,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/magistral-small-1-2-2509": { + "input_cost_per_token": 5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 40000, + "max_output_tokens": 40000, + "max_tokens": 40000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://mistral.ai/pricing#api-pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/mistral-embed": { "input_cost_per_token": 1e-07, "litellm_provider": "mistral", @@ -23238,24 +23475,41 @@ "supports_tool_choice": true }, "mistral/mistral-large-latest": { - "input_cost_per_token": 2e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "mistral", - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 1.5e-06, + "source": "https://docs.mistral.ai/models/mistral-large-3-25-12", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-large-3": { "input_cost_per_token": 5e-07, "litellm_provider": "mistral", - "max_input_tokens": 256000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.mistral.ai/models/mistral-large-3-25-12", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-large-2512": { + "input_cost_per_token": 5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://docs.mistral.ai/models/mistral-large-3-25-12", @@ -23306,14 +23560,30 @@ "input_cost_per_token": 4e-07, "litellm_provider": "mistral", "max_input_tokens": 131072, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-06, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-3-1-2508": { + "input_cost_per_token": 4e-07, + "litellm_provider": "mistral", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/mistral-medium-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-small": { "input_cost_per_token": 1e-07, @@ -23329,17 +23599,79 @@ "supports_tool_choice": true }, "mistral/mistral-small-latest": { - "input_cost_per_token": 1e-07, + "input_cost_per_token": 6e-08, "litellm_provider": "mistral", - "max_input_tokens": 32000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-07, + "output_cost_per_token": 1.8e-07, + "source": "https://mistral.ai/pricing", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-small-3-2-2506": { + "input_cost_per_token": 6e-08, + "litellm_provider": "mistral", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.8e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/ministral-3-3b-2512": { + "input_cost_per_token": 1e-07, + "litellm_provider": "mistral", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/ministral-3-8b-2512": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/ministral-3-14b-2512": { + "input_cost_per_token": 2e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-tiny": { "input_cost_per_token": 2.5e-07, @@ -25657,7 +25989,7 @@ "supports_tool_choice": true }, "openrouter/google/gemini-2.0-flash-001": { - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", @@ -29554,7 +29886,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "us.anthropic.claude-3-5-sonnet-20241022-v2:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -29607,7 +29941,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 2.5e-08, + "cache_creation_input_token_cost": 3.125e-07 }, "us.anthropic.claude-3-opus-20240229-v1:0": { "input_cost_per_token": 1.5e-05, @@ -29620,7 +29956,9 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.5e-06, + "cache_creation_input_token_cost": 1.875e-05 }, "us.anthropic.claude-3-sonnet-20240229-v1:0": { "input_cost_per_token": 3e-06, @@ -29634,7 +29972,9 @@ "supports_pdf_input": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 3e-07, + "cache_creation_input_token_cost": 3.75e-06 }, "us.anthropic.claude-opus-4-1-20250805-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -30527,7 +30867,7 @@ "supports_tool_choice": true }, "vercel_ai_gateway/google/gemini-2.0-flash": { - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_token": 1.5e-07, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, @@ -30541,7 +30881,7 @@ "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.0-flash-lite": { - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_token": 7.5e-08, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, @@ -32059,6 +32399,54 @@ "output_cost_per_token": 3e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-09, + "input_cost_per_audio_token": 2.5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true + }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -37898,7 +38286,7 @@ }, "gemini/gemini-2.0-flash-lite-001": { "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2026-03-31", + "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, "litellm_provider": "gemini", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4934f11d456..42b4f0f7762 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14334,6 +14334,53 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-09, + "input_cost_per_audio_token": 2.5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true + }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -17090,6 +17137,56 @@ "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "supports_service_tier": true }, + "gemini/gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-09, + "input_cost_per_audio_token": 2.5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 250000 + }, "gemini/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, "input_cost_per_audio_token": 1e-06, @@ -32302,6 +32399,54 @@ "output_cost_per_token": 3e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-preview": { + "cache_read_input_token_cost": 2.5e-09, + "input_cost_per_audio_token": 2.5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true + }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, 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 7e8848be301..2e033b6f068 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 @@ -33,8 +33,8 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm.litellm_core_utils.llm_cost_calc.utils import ( - _calculate_input_cost, PromptTokensDetailsResult, + _calculate_input_cost, calculate_cache_writing_cost, generic_cost_per_token, ) @@ -127,6 +127,52 @@ def test_reasoning_tokens_gemini(): ) +def test_reasoning_tokens_gemini_3_1_flash_lite(): + """Test cost calculation for gemini-3.1-flash-lite-preview with reasoning tokens""" + model = "gemini-3.1-flash-lite-preview" + custom_llm_provider = "gemini" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + usage = Usage( + completion_tokens=1000, + prompt_tokens=500, + total_tokens=1500, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=400, + rejected_prediction_tokens=None, + text_tokens=600, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=500, image_tokens=None + ), + ) + model_cost_map = litellm.model_cost[model] + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + + assert round(prompt_cost, 10) == round( + model_cost_map["input_cost_per_token"] * usage.prompt_tokens, + 10, + ) + assert round(completion_cost, 10) == round( + ( + model_cost_map["output_cost_per_token"] + * usage.completion_tokens_details.text_tokens + ) + + ( + model_cost_map["output_cost_per_reasoning_token"] + * usage.completion_tokens_details.reasoning_tokens + ), + 10, + ) + + def test_image_tokens_with_custom_pricing(): """Test that image_tokens in completion are properly costed with output_cost_per_image_token.""" from unittest.mock import patch 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 196bb00f40d..596897f7157 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 @@ -2453,8 +2453,8 @@ def test_gemini_3_image_models_no_thinking_config(): def test_gemini_3_text_models_get_thinking_config(): """ - Test that Gemini 3 text models DO receive automatic thinkingConfig. - This ensures we didn't break the existing behavior for non-image models. + Test that Gemini 3 text models do NOT receive automatic thinkingConfig + when no reasoning_effort or thinking param is provided. """ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -2462,7 +2462,7 @@ def test_gemini_3_text_models_get_thinking_config(): v = VertexGeminiConfig() - # Test gemini-3-pro-preview (text model, should get thinking) + # Test gemini-3-pro-preview (text model, no explicit thinking params) model = "gemini-3-pro-preview" optional_params = {} non_default_params = {} @@ -2474,9 +2474,8 @@ def test_gemini_3_text_models_get_thinking_config(): drop_params=False, ) - # Should have thinkingConfig automatically added - assert "thinkingConfig" in result - assert result["thinkingConfig"]["thinkingLevel"] == "low" + # Should NOT have thinkingConfig automatically added when user provides no reasoning_effort + assert "thinkingConfig" not in result assert result["temperature"] == 1.0 From deb8fea6b114eee809745c65995249f6d6b6457f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 15:19:16 +0530 Subject: [PATCH 036/147] Add blog post for gemini-3.1-flash-lite-preview --- .../blog/gemini_3_1_flash_lite/index.md | 175 ++++++++++++++++++ 1 file changed, 175 insertions(+) create mode 100644 docs/my-website/blog/gemini_3_1_flash_lite/index.md diff --git a/docs/my-website/blog/gemini_3_1_flash_lite/index.md b/docs/my-website/blog/gemini_3_1_flash_lite/index.md new file mode 100644 index 00000000000..9ef4bacb2ad --- /dev/null +++ b/docs/my-website/blog/gemini_3_1_flash_lite/index.md @@ -0,0 +1,175 @@ +--- +slug: gemini_3_1_flash_lite_preview +title: "DAY 0 Support: Gemini 3.1 Flash Lite Preview on LiteLLM" +date: 2026-03-03T08:00:00 +authors: + - name: Sameer Kankute + title: SWE @ LiteLLM (LLM Translation) + url: https://www.linkedin.com/in/sameer-kankute/ + image_url: https://pbs.twimg.com/profile_images/2001352686994907136/ONgNuSk5_400x400.jpg + - name: Krrish Dholakia + title: "CEO, LiteLLM" + url: https://www.linkedin.com/in/krish-d/ + image_url: https://pbs.twimg.com/profile_images/1298587542745358340/DZv3Oj-h_400x400.jpg + - name: Ishaan Jaff + title: "CTO, LiteLLM" + url: https://www.linkedin.com/in/reffajnaahsi/ + image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg +description: "Guide to using Gemini 3.1 Flash Lite Preview on LiteLLM Proxy and SDK with day 0 support." +tags: [gemini, day 0 support, llms, supernova] +hide_table_of_contents: false +--- + + +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Gemini 3.1 Flash Lite Preview Day 0 Support + +LiteLLM now supports `gemini-3.1-flash-lite-preview` with full day 0 support! + +:::note +If you only want cost tracking, you need no change in your current Litellm version. But if you want the support for new features introduced along with it like thinking levels, you will need to use v1.80.8-stable.1 or above. +::: + +## Deploy this version + + + + +``` showLineNumbers title="docker run litellm" +docker run \ +-e STORE_MODEL_IN_DB=True \ +-p 4000:4000 \ +ghcr.io/berriai/litellm:main-v1.80.8-stable.1 +``` + + + + + +``` showLineNumbers title="pip install litellm" +pip install litellm==v1.80.8-stable.1 +``` + + + + +## What's New + +Supports all four thinking levels: +- **MINIMAL**: Ultra-fast responses with minimal reasoning +- **LOW**: Simple instruction following +- **MEDIUM**: Balanced reasoning for complex tasks +- **HIGH**: Maximum reasoning depth (dynamic) + +--- + +## Quick Start + + + + +**Basic Usage** + +```python +from litellm import completion + +response = completion( + model="gemini/gemini-3.1-flash-lite-preview", + messages=[{"role": "user", "content": "Extract key entities from this text: ..."}], +) + +print(response.choices[0].message.content) +``` + +**With Thinking Levels** + +```python +from litellm import completion + +# Use MEDIUM thinking for complex reasoning tasks +response = completion( + model="gemini/gemini-3.1-flash-lite-preview", + messages=[{"role": "user", "content": "Analyze this dataset and identify patterns"}], + reasoning_effort="medium", # low, medium , high +) + +print(response.choices[0].message.content) +``` + + + + + +**1. Setup config.yaml** + +```yaml +model_list: + - model_name: gemini-3.1-flash-lite + litellm_params: + model: gemini/gemini-3.1-flash-lite-preview + api_key: os.environ/GEMINI_API_KEY + + # Or use Vertex AI + - model_name: vertex-gemini-3.1-flash-lite + litellm_params: + model: vertex_ai/gemini-3.1-flash-lite-preview + vertex_project: your-project-id + vertex_location: us-central1 +``` + +**2. Start proxy** + +```bash +litellm --config /path/to/config.yaml +``` + +**3. Make requests** + +```bash +curl -X POST http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{ + "model": "gemini-3.1-flash-lite", + "messages": [{"role": "user", "content": "Extract structured data from this text"}], + "reasoning_effort": "low" + }' +``` + + + + +--- + +## Supported Endpoints + +LiteLLM provides **full end-to-end support** for Gemini 3.1 Flash Lite Preview on: + +- ✅ `/v1/chat/completions` - OpenAI-compatible chat completions endpoint +- ✅ `/v1/responses` - OpenAI Responses API endpoint (streaming and non-streaming) +- ✅ [`/v1/messages`](../../docs/anthropic_unified) - Anthropic-compatible messages endpoint +- ✅ `/v1/generateContent` – [Google Gemini API](../../docs/generateContent.md) compatible endpoint + +All endpoints support: +- Streaming and non-streaming responses +- Function calling with thought signatures +- Multi-turn conversations +- All Gemini 3-specific features (thinking levels, thought signatures) +- Full multimodal support (text, image, audio, video) + +--- + +## `reasoning_effort` Mapping for Gemini 3.1 + +LiteLLM automatically maps OpenAI's `reasoning_effort` parameter to Gemini's `thinkingLevel`: + +| reasoning_effort | thinking_level | Use Case | +|------------------|----------------|----------| +| `minimal` | `minimal` | Ultra-fast responses, simple queries | +| `low` | `low` | Basic instruction following | +| `medium` | `medium` | Balanced reasoning for moderate complexity | +| `high` | `high` | Maximum reasoning depth, complex problems | +| `disable` | `minimal` | Disable extended reasoning | +| `none` | `minimal` | No extended reasoning | \ No newline at end of file From 409208771e6fa8bb596f580e6efd7bcb36b0a308 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 09:24:49 -0300 Subject: [PATCH 037/147] fix(pricing): add 5 missing OpenRouter model pricing entries Fixes #22609 Adds pricing for OpenRouter models that were routing correctly but returning $0 for spend tracking due to missing cost map entries: - openrouter/anthropic/claude-sonnet-4.6 ($3.00/$15.00 per 1M tokens) - openrouter/google/gemini-3.1-pro-preview ($2.00/$12.00 per 1M tokens) - openrouter/openai/gpt-5.1-codex-max ($1.25/$10.00 per 1M tokens) - openrouter/qwen/qwen3-coder-plus ($1.00/$5.00 per 1M tokens) - openrouter/z-ai/glm-5 ($0.80/$2.56 per 1M tokens) --- model_prices_and_context_window.json | 106 +++++++++++++++++++++++++++ 1 file changed, 106 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4c0694db371..07b44fb7809 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25373,6 +25373,30 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "openrouter/anthropic/claude-sonnet-4.6": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "source": "https://openrouter.ai/anthropic/claude-sonnet-4.6", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 159 + }, "openrouter/anthropic/claude-opus-4.5": { "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, @@ -25723,6 +25747,39 @@ "supports_web_search": true, "tpm": 800000 }, + "openrouter/google/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "source": "https://openrouter.ai/google/gemini-3.1-pro-preview", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 1.875e-06, "litellm_provider": "openrouter", @@ -26100,6 +26157,29 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/openai/gpt-5.1-codex-max": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://openrouter.ai/openai/gpt-5.1-codex-max", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "openrouter/openai/gpt-5.2": { "input_cost_per_image": 0, "cache_read_input_token_cost": 1.75e-07, @@ -26254,6 +26334,19 @@ "supports_tool_choice": true, "supports_function_calling": true }, + "openrouter/qwen/qwen3-coder-plus": { + "input_cost_per_token": 1e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://openrouter.ai/qwen/qwen3-coder-plus", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/qwen/qwen3-235b-a22b-2507": { "input_cost_per_token": 7.1e-08, "litellm_provider": "openrouter", @@ -26389,6 +26482,19 @@ "supports_vision": true, "supports_prompt_caching": false }, + "openrouter/z-ai/glm-5": { + "input_cost_per_token": 8e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 202752, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.56e-06, + "source": "https://openrouter.ai/z-ai/glm-5", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/minimax/minimax-m2.1": { "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1.2e-06, From 79771261819774fbe40c4a4201148b872479b5a3 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 09:52:17 -0300 Subject: [PATCH 038/147] fix(pricing): add 7 missing DashScope model pricing entries Fixes #22646 Adds pricing for DashScope models that were missing from the cost map, causing $0 spend tracking in the proxy dashboard: - dashscope/qwen3-max-2026-01-23 (tiered, same as qwen3-max) - dashscope/qwen3-next-80b-a3b-instruct ($0.15/$1.20 per 1M) - dashscope/qwen3-next-80b-a3b-thinking ($0.15/$1.20 per 1M) - dashscope/qwen3-vl-235b-a22b-instruct ($0.40/$1.60 per 1M) - dashscope/qwen3-vl-235b-a22b-thinking ($0.40/$4.00 per 1M) - dashscope/qwen3-vl-32b-instruct ($0.16/$0.64 per 1M) - dashscope/qwen3-vl-32b-thinking ($0.16/$2.87 per 1M) --- model_prices_and_context_window.json | 116 +++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4c0694db371..7f59fcb0818 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9779,6 +9779,122 @@ } ] }, + "dashscope/qwen3-max-2026-01-23": { + "litellm_provider": "dashscope", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "dashscope/qwen3-next-80b-a3b-instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "dashscope/qwen3-next-80b-a3b-thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "dashscope/qwen3-vl-235b-a22b-instruct": { + "input_cost_per_token": 4e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "dashscope/qwen3-vl-235b-a22b-thinking": { + "input_cost_per_token": 4e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "dashscope/qwen3-vl-32b-instruct": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.4e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "dashscope/qwen3-vl-32b-thinking": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.87e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "dashscope/qwen3-vl-plus": { "litellm_provider": "dashscope", "max_input_tokens": 260096, From 058fac848e0f7e16ea0985e9ed7eae897523bb60 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 16:20:51 +0530 Subject: [PATCH 039/147] Add Encrypted-content-aware deployment affinity for the Router --- .../encrypted_content_affinity_check.py | 267 ++++++++++++++++++ 1 file changed, 267 insertions(+) create mode 100644 litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py 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 new file mode 100644 index 00000000000..cdf69e4b24c --- /dev/null +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -0,0 +1,267 @@ +""" +Encrypted-content-aware deployment affinity for the Router. + +When Codex or other models use `store: false` with `include: ["reasoning.encrypted_content"]`, +the response output items contain encrypted reasoning tokens tied to the originating +organization's API key. If a follow-up request containing those items is routed to a +different deployment (different org), OpenAI rejects it with an `invalid_encrypted_content` +error because the organization_id doesn't match. + +This callback solves the problem by: +1. Tracking output item IDs from Responses API responses and mapping them to the + deployment (model_id) that produced them. +2. On subsequent requests, scanning the `input` field for known item IDs and pinning + the request to the originating deployment. + +Safe to enable globally: +- Only activates when known item IDs appear in the request `input`. +- No effect on embedding models, chat completions, or first-time requests. +- No quota reduction -- first requests are fully load balanced. +""" + +from typing import Any, List, Optional, cast + +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse + +_DEFAULT_TTL_SECONDS = 86400 # 24 hours + + +class EncryptedContentAffinityCheck(CustomLogger): + """ + Routes follow-up Responses API requests to the deployment that produced + the encrypted output items they reference. + + Wired via ``Router(optional_pre_call_checks=["encrypted_content_affinity"])``. + """ + + CACHE_KEY_PREFIX = "encrypted_content_affinity:v1" + + def __init__( + self, + cache: DualCache, + ttl_seconds: int = _DEFAULT_TTL_SECONDS, + ): + super().__init__() + self.cache = cache + self.ttl_seconds = ttl_seconds + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def _get_output_from_response( + response_obj: Any, + ) -> Optional[list]: + """ + Extract the ``output`` list from a Responses API response, handling + both ``ResponsesAPIResponse`` objects and plain dicts. + """ + if isinstance(response_obj, ResponsesAPIResponse): + return response_obj.output + if isinstance(response_obj, dict) and "output" in response_obj: + output = response_obj["output"] + if isinstance(output, list): + return output + if hasattr(response_obj, "output"): + output = response_obj.output + if isinstance(output, list): + return output + return None + + @staticmethod + def _extract_item_ids_from_output( + output: list, + ) -> List[str]: + """Extract all item IDs from a Responses API output list.""" + item_ids: List[str] = [] + for item in output: + item_id: Optional[str] = None + if isinstance(item, dict): + item_id = item.get("id") + else: + item_id = getattr(item, "id", None) + if item_id and isinstance(item_id, str): + item_ids.append(item_id) + return item_ids + + @staticmethod + def _extract_item_ids_from_input(request_input: Any) -> List[str]: + """ + Extract item IDs from the ``input`` field of a Responses API request. + + ``input`` can be: + - a plain string -> no item IDs + - a list of items -> each item may have an ``id`` field + """ + if not isinstance(request_input, list): + return [] + + item_ids: List[str] = [] + for item in request_input: + if isinstance(item, dict): + item_id = item.get("id") + if item_id and isinstance(item_id, str): + item_ids.append(item_id) + return item_ids + + @classmethod + def _cache_key(cls, item_id: str) -> str: + return f"{cls.CACHE_KEY_PREFIX}:{item_id}" + + @staticmethod + def _find_deployment_by_model_id( + healthy_deployments: List[dict], model_id: str + ) -> Optional[dict]: + for deployment in healthy_deployments: + model_info = deployment.get("model_info") + if not isinstance(model_info, dict): + continue + deployment_model_id = model_info.get("id") + if deployment_model_id is not None and str(deployment_model_id) == str( + model_id + ): + return deployment + return None + + @staticmethod + def _get_model_id_from_kwargs(kwargs: dict) -> Optional[str]: + """ + Extract the deployment model_id from success-callback kwargs. + + The Router populates ``litellm_params.metadata.model_info.id`` after + selecting a deployment. Also check top-level ``model_info`` as a + fallback (some call paths set it there). + """ + # Primary path: litellm_params -> metadata -> model_info -> id + litellm_params = kwargs.get("litellm_params") + if isinstance(litellm_params, dict): + metadata = litellm_params.get("metadata") + if isinstance(metadata, dict): + model_info = metadata.get("model_info") + if isinstance(model_info, dict): + model_id = model_info.get("id") + if model_id is not None: + return str(model_id) + + # Fallback: top-level model_info (set by some router call paths) + model_info = kwargs.get("model_info") + if isinstance(model_info, dict): + model_id = model_info.get("id") + if model_id is not None: + return str(model_id) + + return None + + # ------------------------------------------------------------------ + # Response tracking (success callback) + # ------------------------------------------------------------------ + + async def async_log_success_event( + self, kwargs: dict, response_obj: Any, start_time: Any, end_time: Any + ) -> None: + """ + After a successful Responses API call, cache each output item ID + mapped to the deployment that produced it. + """ + output = self._get_output_from_response(response_obj) + if output is None: + return + + model_id = self._get_model_id_from_kwargs(kwargs) + if not model_id: + verbose_router_logger.debug( + "EncryptedContentAffinityCheck: model_id not found in kwargs, skipping tracking", + ) + return + + item_ids = self._extract_item_ids_from_output(output) + if not item_ids: + return + + for item_id in item_ids: + try: + cache_key = self._cache_key(item_id) + await self.cache.async_set_cache( + cache_key, + model_id, + ttl=self.ttl_seconds, + ) + except Exception as e: + verbose_router_logger.debug( + "EncryptedContentAffinityCheck: failed to cache item_id=%s error=%s", + item_id, + e, + ) + + verbose_router_logger.info( + "EncryptedContentAffinityCheck: cached %d item IDs -> deployment=%s", + len(item_ids), + model_id, + ) + + # ------------------------------------------------------------------ + # Request routing (pre-call filter) + # ------------------------------------------------------------------ + + async def async_filter_deployments( + self, + model: str, + healthy_deployments: List, + messages: Optional[List[AllMessageValues]], + request_kwargs: Optional[dict] = None, + parent_otel_span: Optional[Span] = None, + ) -> List[dict]: + """ + If the request ``input`` contains items whose IDs were previously + tracked, pin the request to the deployment that produced them. + """ + request_kwargs = request_kwargs or {} + typed_healthy_deployments = cast(List[dict], healthy_deployments) + + request_input = request_kwargs.get("input") + input_item_ids = self._extract_item_ids_from_input(request_input) + if not input_item_ids: + return typed_healthy_deployments + + verbose_router_logger.debug( + "EncryptedContentAffinityCheck: found %d item IDs in input, checking cache", + len(input_item_ids), + ) + + for item_id in input_item_ids: + cache_key = self._cache_key(item_id) + try: + cached_model_id = await self.cache.async_get_cache(key=cache_key) + except Exception: + continue + + if not cached_model_id or not isinstance(cached_model_id, str): + continue + + deployment = self._find_deployment_by_model_id( + healthy_deployments=typed_healthy_deployments, + model_id=cached_model_id, + ) + if deployment is not None: + verbose_router_logger.info( + "EncryptedContentAffinityCheck: item_id=%s pinning -> deployment=%s", + item_id, + cached_model_id, + ) + request_kwargs[ + "_encrypted_content_affinity_pinned" + ] = True + return [deployment] + + verbose_router_logger.info( + "EncryptedContentAffinityCheck: cached deployment=%s for item_id=%s " + "not found in healthy_deployments", + cached_model_id, + item_id, + ) + + return typed_healthy_deployments From 92b255628216fe6550928eaa53fbf18b7780c519 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 16:21:20 +0530 Subject: [PATCH 040/147] Add encrypted_content_affinity in router --- litellm/router.py | 30 ++++++++++++++++++++++++++++++ litellm/types/router.py | 1 + 2 files changed, 31 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index 5c8d27e76d7..7dee4fa83de 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -115,6 +115,9 @@ from litellm.router_utils.handle_error import ( 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, +) from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) @@ -1248,6 +1251,25 @@ class Router: self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + # --------------------------------------------------------------------- + # Encrypted content affinity + # --------------------------------------------------------------------- + if "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( + cache=self.cache, + ttl_seconds=self.deployment_affinity_ttl_seconds, + ) + self.optional_callbacks.append(ec_callback) + litellm.logging_callback_manager.add_litellm_callback(ec_callback) + # --------------------------------------------------------------------- # Remaining optional pre-call checks # --------------------------------------------------------------------- @@ -1257,6 +1279,7 @@ class Router: "deployment_affinity", "responses_api_deployment_check", "session_affinity", + "encrypted_content_affinity", ): continue if pre_call_check == "prompt_caching": @@ -8808,6 +8831,13 @@ class Router: if isinstance(healthy_deployments, dict): return healthy_deployments + # When encrypted content affinity pins to a specific deployment, + if ( + request_kwargs.get("_encrypted_content_affinity_pinned") + and len(healthy_deployments) == 1 + ): + return healthy_deployments[0] + start_time = time.time() if ( self.routing_strategy == "usage-based-routing-v2" diff --git a/litellm/types/router.py b/litellm/types/router.py index aa4d7bd9a97..fca731d1f91 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -813,6 +813,7 @@ OptionalPreCallChecks = List[ "session_affinity", "forward_client_headers_by_model_group", "enforce_model_rate_limits", + "encrypted_content_affinity", ] ] From 394c49d3037a9cc05cc942cc73c13a9f7acb4eb6 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 16:21:49 +0530 Subject: [PATCH 041/147] Add tests for encrypted_content_affinity --- .../test_encrypted_content_affinity_check.py | 335 ++++++++++++++++++ 1 file changed, 335 insertions(+) create mode 100644 tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py 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 new file mode 100644 index 00000000000..a266df749c1 --- /dev/null +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -0,0 +1,335 @@ +import asyncio +import os +import sys +from unittest.mock import AsyncMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import json + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, +) + + +class MockResponse: + def __init__(self, json_data, status_code): + self._json_data = json_data + self.status_code = status_code + self.text = json.dumps(json_data) + self.headers = {} + + def json(self): + return self._json_data + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_tracks_and_routes(): + """ + When encrypted_content_affinity is enabled, output item IDs from responses + are tracked, and follow-up requests containing those IDs route to the same + deployment. + """ + mock_response_data = { + "id": "resp_mock-123", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "openai/gpt-5.1-codex", + "output": [ + { + "type": "message", + "id": "msg_abc123", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], + }, + { + "type": "reasoning", + "id": "rs_encrypted_item_456", + "status": "completed", + }, + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 5, + "output_tokens": 10, + "total_tokens": 15, + }, + "error": None, + } + + router = litellm.Router( + model_list=[ + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-1", + }, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-2", + }, + "model_info": {"id": "deployment-2"}, + }, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + ) + + model_group = "openai.gpt-5.1-codex" + + # Track which deployment was selected + selected_deployments = [] + + def deterministic_choice(seq): + # First call: select deployment-1 + # Second call: would select deployment-2, but affinity should override + if len(selected_deployments) == 0: + return seq[0] + return seq[1] if len(seq) > 1 else seq[0] + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, patch( + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=deterministic_choice, + ): + mock_post.return_value = MockResponse(mock_response_data, 200) + + # First request: no encrypted items in input + first_response = await router.aresponses( + model=model_group, + input="Hello, how are you?", + ) + first_model_id = first_response._hidden_params["model_id"] + selected_deployments.append(first_model_id) + + # Give async callbacks time to run + await asyncio.sleep(0.2) + + # Second request: includes encrypted item IDs from first response + second_response = await router.aresponses( + model=model_group, + input=[ + {"type": "message", "id": "msg_abc123", "role": "assistant"}, + {"type": "reasoning", "id": "rs_encrypted_item_456"}, + ], + ) + second_model_id = second_response._hidden_params["model_id"] + + # Affinity should route to the same deployment + assert second_model_id == first_model_id, ( + f"Expected affinity to route to {first_model_id}, " + f"but got {second_model_id}" + ) + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_no_effect_on_chat_completions(): + """ + Encrypted content affinity should not affect regular chat completions + (they don't use the Responses API). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": os.environ.get("OPENAI_API_KEY", "test-key"), + }, + "model_info": {"id": "chat-deployment-1"}, + }, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + ) + + mock_chat_response = { + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-3.5-turbo", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}, + } + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = MockResponse(mock_chat_response, 200) + + # Multiple chat completion requests should work normally + response1 = await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello"}], + ) + response2 = await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello again"}], + ) + + # Both should succeed (no affinity interference) + # Check that responses have IDs (litellm may modify them) + assert response1.id is not None + assert response2.id is not None + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_bypasses_rpm_limits(): + """ + When encrypted content affinity pins to a deployment, it should bypass + RPM limits since the encrypted content will fail on any other deployment. + """ + mock_response_data = { + "id": "resp_mock-rpm-test", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "openai/gpt-5.1-codex", + "output": [ + { + "type": "reasoning", + "id": "rs_encrypted_must_pin", + "status": "completed", + }, + ], + "usage": {"input_tokens": 5, "output_tokens": 10, "total_tokens": 15}, + "error": None, + } + + router = litellm.Router( + model_list=[ + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-1", + "rpm": 1, # Very low limit + }, + "model_info": {"id": "rpm-limited-deployment"}, + }, + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-2", + "rpm": 100, + }, + "model_info": {"id": "high-rpm-deployment"}, + }, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + routing_strategy="usage-based-routing-v2", + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = MockResponse(mock_response_data, 200) + + # First request goes to the low-RPM deployment + first_response = await router.aresponses( + model="openai.gpt-5.1-codex", + input="Initial request", + ) + first_model_id = first_response._hidden_params["model_id"] + + await asyncio.sleep(0.2) + + # Second request with encrypted content should pin to the same deployment + # even though it's at RPM limit + second_response = await router.aresponses( + model="openai.gpt-5.1-codex", + input=[ + {"type": "reasoning", "id": "rs_encrypted_must_pin"}, + ], + ) + second_model_id = second_response._hidden_params["model_id"] + + # Should route to the same deployment despite RPM limit + assert second_model_id == first_model_id + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_no_match_normal_routing(): + """ + When input contains item IDs that aren't tracked, normal load balancing + should occur. + """ + mock_response_data = { + "id": "resp_mock-no-match", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "openai/gpt-5.1-codex", + "output": [ + { + "type": "message", + "id": "msg_new", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Response"}], + }, + ], + "usage": {"input_tokens": 5, "output_tokens": 10, "total_tokens": 15}, + "error": None, + } + + router = litellm.Router( + model_list=[ + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-1", + }, + "model_info": {"id": "deployment-a"}, + }, + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-2", + }, + "model_info": {"id": "deployment-b"}, + }, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = MockResponse(mock_response_data, 200) + + # Request with unknown item IDs should use normal routing + response = await router.aresponses( + model="openai.gpt-5.1-codex", + input=[ + {"type": "message", "id": "unknown_item_id_12345"}, + ], + ) + + # Should succeed with normal routing (litellm may modify the ID) + assert response.id is not None + # Verify it contains the original response ID in some form + assert "resp_mock-no-match" in str(response.id) or response.id.startswith("resp_") From 9f627c67d83bf8216d2637ba9eb877e0588a42ca Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 16:22:00 +0530 Subject: [PATCH 042/147] Add tests for encrypted_content_affinity --- .../pre_call_checks/test_encrypted_content_affinity_check.py | 4 ---- 1 file changed, 4 deletions(-) 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 a266df749c1..c70290b674d 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 @@ -10,10 +10,6 @@ sys.path.insert(0, os.path.abspath("../..")) import json import litellm -from litellm.caching.dual_cache import DualCache -from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, -) class MockResponse: From fbec5c5ccf6c3633cae358f50da695203a76a624 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 16:22:15 +0530 Subject: [PATCH 043/147] Add docs for encrypted_content_affinity --- docs/my-website/docs/proxy/config_settings.md | 2 +- docs/my-website/docs/proxy/load_balancing.md | 33 ++++ docs/my-website/docs/response_api.md | 155 ++++++++++++++++++ 3 files changed, 189 insertions(+), 1 deletion(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 7b2011e45dd..af868bc9f9d 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -360,7 +360,7 @@ router_settings: | redis_url | str | URL for Redis server. **Known performance issue with Redis URL.** | | cache_responses | boolean | Flag to enable caching LLM Responses, if cache set under `router_settings`. If true, caches responses. Defaults to False. | | router_general_settings | RouterGeneralSettings | [SDK-Only] Router general settings - contains optimizations like 'async_only_mode'. [Docs](../routing.md#router-general-settings) | -| optional_pre_call_checks | List[str] | List of pre-call checks to add to the router. Supported: `router_budget_limiting`, `prompt_caching`, `responses_api_deployment_check`, `deployment_affinity`, `forward_client_headers_by_model_group` | +| optional_pre_call_checks | List[str] | List of pre-call checks to add to the router. Supported: `router_budget_limiting`, `prompt_caching`, `responses_api_deployment_check`, `encrypted_content_affinity`, `deployment_affinity`, `session_affinity`, `forward_client_headers_by_model_group` | | deployment_affinity_ttl_seconds | int | TTL (seconds) for user-key → deployment affinity mapping when `deployment_affinity` is enabled (configured at Router init / proxy startup). Defaults to `3600` (1 hour). | | ignore_invalid_deployments | boolean | If true, ignores invalid deployments. Default for proxy is True - to prevent invalid models from blocking other models from being loaded. | | search_tools | List[SearchToolTypedDict] | List of search tool configurations for Search API integration. Each tool specifies a search_tool_name and litellm_params with search_provider, api_key, api_base, etc. [Further Docs](../search.md) | diff --git a/docs/my-website/docs/proxy/load_balancing.md b/docs/my-website/docs/proxy/load_balancing.md index 186307d6498..5bf39d179f6 100644 --- a/docs/my-website/docs/proxy/load_balancing.md +++ b/docs/my-website/docs/proxy/load_balancing.md @@ -347,3 +347,36 @@ If `order=1` deployment is unavailable (e.g., rate-limited), the router falls ba - **Higher throughput**: More requests handled simultaneously across deployments - **Improved reliability**: If one deployment fails, traffic automatically routes to healthy ones - **Better resource utilization**: Load spread evenly across all available deployments + +## Special Considerations for Responses API + +When load balancing OpenAI's Responses API across deployments with **different API keys** (e.g., different Azure regions or organizations), encrypted content items (like `rs_...` reasoning items) can only be decrypted by the originating API key. + +**Solution:** Use the `encrypted_content_affinity` pre-call check to automatically route follow-up requests containing encrypted items to the correct deployment: + +```yaml +model_list: + - model_name: gpt-5.1-codex + litellm_params: + model: azure/gpt-5.1-codex + api_base: https://eastus.openai.azure.com/ + api_key: os.environ/AZURE_API_KEY_EASTUS + model_info: + id: "deployment-eastus" + + - model_name: gpt-5.1-codex + litellm_params: + model: azure/gpt-5.1-codex + api_base: https://westeurope.openai.azure.com/ + api_key: os.environ/AZURE_API_KEY_WESTEUROPE + model_info: + id: "deployment-westeurope" + +router_settings: + optional_pre_call_checks: + - encrypted_content_affinity # 👈 Prevents invalid_encrypted_content errors +``` + +This ensures requests containing encrypted content are routed to the deployment that created them, while other requests continue to load balance normally. + +**[Learn more about Encrypted Content Affinity →](../response_api.md#encrypted-content-affinity-multi-region-load-balancing)** diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index b37be2b5bc2..85144a80245 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -920,9 +920,14 @@ follow_up = await router.aresponses( To enable session continuity for Responses API in your LiteLLM proxy, set `optional_pre_call_checks` in your proxy config.yaml. - `responses_api_deployment_check`: high priority routing when `previous_response_id` is provided +- `encrypted_content_affinity`: **[Recommended]** content-aware routing for encrypted items (e.g., `rs_...` reasoning items) - `session_affinity`: sticky sessions based on session id (takes priority over `deployment_affinity`) - `deployment_affinity`: sticky sessions based on user key (applies even without `previous_response_id`) +:::tip Recommended: Use `encrypted_content_affinity` +For Responses API with load balancing across deployments with **different API keys**, use `encrypted_content_affinity` instead of `deployment_affinity`. It only pins requests that contain encrypted content, avoiding quota reduction while preventing `invalid_encrypted_content` errors. +::: + Notes: - User-key affinity is keyed on `metadata.user_api_key_hash` (the API key hash). The OpenAI `user` request parameter is an end-user identifier and is intentionally not used for deployment affinity. - Session-ID affinity is keyed on `metadata.session_id`. For proxy requests, this can be passed via the `x-litellm-session-id` HTTP header. For Python SDK requests, you can pass it via `litellm_metadata={"session_id": "value"}` in request args. @@ -983,6 +988,156 @@ follow_up = client.responses.create( +## Encrypted Content Affinity (Multi-Region Load Balancing) + +When load balancing Responses API across deployments with **different API keys** (e.g., different Azure regions or OpenAI organizations), encrypted content items (like `rs_...` reasoning items) can only be decrypted by the API key that created them. + +### The Problem + +```json +{ + "error": { + "message": "The encrypted content for item rs_0d09d6e56879e76500699d6feee41c8197bd268aae76141f87 could not be verified. Reason: Encrypted content organization_id did not match the target organization.", + "type": "invalid_request_error", + "code": "invalid_encrypted_content" + } +} +``` + +This error occurs when: +1. Initial request goes to Deployment A (API Key 1) → produces encrypted item `rs_xyz` +2. Follow-up request with `rs_xyz` in input gets load balanced to Deployment B (API Key 2) +3. Deployment B cannot decrypt content created by Deployment A → **request fails** + +### The Solution: `encrypted_content_affinity` + +The `encrypted_content_affinity` pre-call check intelligently tracks encrypted content and routes follow-up requests to the originating deployment **only when necessary**. + +**Key Benefits:** +- ✅ **No quota reduction**: Unlike `deployment_affinity`, only pins requests that contain tracked encrypted items +- ✅ **Bypasses rate limits**: When encrypted content requires a specific deployment, RPM/TPM limits are bypassed (the request would fail on any other deployment anyway) +- ✅ **No `previous_response_id` required**: Works by tracking item IDs in response output and matching them in request input +- ✅ **Globally safe**: Can be enabled for all models; non-Responses-API calls (chat, embeddings) are unaffected + +### How It Works + +1. **Tracking Phase** (after successful response): + - Extracts all item IDs from response `output` (e.g., `msg_abc`, `rs_xyz`) + - Caches mapping: `item_id` → `deployment_id` (default TTL: 24 hours) + +2. **Routing Phase** (before request): + - Scans request `input` for item IDs + - If tracked item found → pins to originating deployment, bypasses rate limits + - If no tracked items → normal load balancing + +### Configuration + + + + +```python +from litellm import Router + +router = Router( + model_list=[ + { + "model_name": "gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "org-1-api-key", # Different API key + }, + "model_info": {"id": "deployment-us-east"}, + }, + { + "model_name": "gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "org-2-api-key", # Different API key + }, + "model_info": {"id": "deployment-eu-west"}, + }, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + deployment_affinity_ttl_seconds=86400, # 24 hours (default) +) + +# Initial request - routes to any deployment +response1 = await router.aresponses( + model="gpt-5.1-codex", + input="Explain quantum computing", +) + +# Follow-up with encrypted items - automatically routes to same deployment +response2 = await router.aresponses( + model="gpt-5.1-codex", + input=response1.output, # Contains encrypted items from response1 +) +``` + + + + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: gpt-5.1-codex + litellm_params: + model: azure/gpt-5.1-codex + api_base: https://eastus.openai.azure.com/ + api_key: os.environ/AZURE_API_KEY_EASTUS + rpm: 600 + tpm: 100000 + model_info: + id: "gpt-5.1-codex-eastus" + + - model_name: gpt-5.1-codex + litellm_params: + model: azure/gpt-5.1-codex + api_base: https://westeurope.openai.azure.com/ + api_key: os.environ/AZURE_API_KEY_WESTEUROPE + rpm: 600 + tpm: 100000 + model_info: + id: "gpt-5.1-codex-westeurope" + +router_settings: + routing_strategy: usage-based-routing-v2 + enable_pre_call_checks: true + optional_pre_call_checks: + - encrypted_content_affinity + deployment_affinity_ttl_seconds: 86400 # Optional, default is 86400 (24 hours) +``` + +**Start proxy:** +```bash +litellm --config config.yaml +``` + + + + +### When to Use Each Affinity Type + +| Affinity Type | Use Case | Scope | Quota Impact | +|---------------|----------|-------|--------------| +| **`encrypted_content_affinity`** | **[Recommended]** Multi-region Responses API with different API keys | Only requests with tracked encrypted items | ✅ None (surgical pinning) | +| `responses_api_deployment_check` | When `previous_response_id` is available | Requests with `previous_response_id` | ✅ None | +| `session_affinity` | Session-based applications | All requests with same `session_id` | ⚠️ Reduces quota by # of sessions | +| `deployment_affinity` | Simple sticky sessions | All requests from same API key | ❌ Reduces quota by # of users | + +### Multi-Instance Deployment (Redis) + +For multiple LiteLLM proxy instances, use Redis to share affinity state: + +```yaml +router_settings: + optional_pre_call_checks: + - encrypted_content_affinity + redis_host: redis.example.com + redis_port: 6379 + redis_password: your-password +``` + + ## Calling non-Responses API endpoints (`/responses` to `/chat/completions` Bridge) LiteLLM allows you to call non-Responses API models via a bridge to LiteLLM's `/chat/completions` endpoint. This is useful for calling Anthropic, Gemini and even non-Responses API OpenAI models. From 37612bdf56349635b1c7d97066d9ee66988fb4d6 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 16:24:57 +0530 Subject: [PATCH 044/147] ADd incident report --- .../index.md | 231 ++++++++++++++++++ 1 file changed, 231 insertions(+) create mode 100644 docs/my-website/blog/responses_api_encrypted_content_incident/index.md diff --git a/docs/my-website/blog/responses_api_encrypted_content_incident/index.md b/docs/my-website/blog/responses_api_encrypted_content_incident/index.md new file mode 100644 index 00000000000..8422f6d3cd0 --- /dev/null +++ b/docs/my-website/blog/responses_api_encrypted_content_incident/index.md @@ -0,0 +1,231 @@ +--- +slug: responses-api-encrypted-content-incident +title: "Incident Report: Encrypted Content Failures in Multi-Region Responses API Load Balancing" +date: 2026-02-24T10:00:00 +authors: + - name: Sameer Kankute + title: SWE @ LiteLLM (LLM Translation) + url: https://www.linkedin.com/in/sameer-kankute/ + image_url: https://pbs.twimg.com/profile_images/2001352686994907136/ONgNuSk5_400x400.jpg + - name: Krrish Dholakia + title: "CEO, LiteLLM" + url: https://www.linkedin.com/in/krish-d/ + image_url: https://pbs.twimg.com/profile_images/1298587542745358340/DZv3Oj-h_400x400.jpg + - name: Ishaan Jaff + title: "CTO, LiteLLM" + url: https://www.linkedin.com/in/reffajnaahsi/ + image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg +tags: [incident-report, proxy, responses-api, load-balancing] +hide_table_of_contents: false +--- + +**Date:** Feb 24, 2026 +**Duration:** Ongoing (until fix deployed) +**Severity:** High (for users load balancing Responses API across different API keys) +**Status:** Resolved + +## Summary + +When load balancing OpenAI's Responses API across deployments with **different API keys** (e.g., different Azure regions or OpenAI organizations), follow-up requests containing encrypted content items (like `rs_...` reasoning items) would fail with: + +```json +{ + "error": { + "message": "The encrypted content for item rs_0d09d6e56879e76500699d6feee41c8197bd268aae76141f87 could not be verified. Reason: Encrypted content organization_id did not match the target organization.", + "type": "invalid_request_error", + "code": "invalid_encrypted_content" + } +} +``` + +Encrypted content items are cryptographically tied to the API key's organization that created them. When the router load balanced a follow-up request to a deployment with a different API key, decryption failed. + +- **Responses API calls with encrypted content:** Complete failure when routed to wrong deployment +- **Initial requests:** Unaffected — only follow-up requests containing encrypted items failed +- **Other API endpoints:** No impact — chat completions, embeddings, etc. functioned normally + +{/* truncate */} + +--- + +## Background + +OpenAI's Responses API can return encrypted "reasoning items" (with IDs like `rs_...`) that contain intermediate reasoning steps. These items are encrypted with the organization's key and can only be decrypted by the same organization's API key. + +When load balancing across deployments with different API keys, the existing affinity mechanisms were insufficient: + +- **`responses_api_deployment_check`**: Requires `previous_response_id` which some clients (like Codex) don't provide +- **`deployment_affinity`**: Too broad — pins *all* requests from a user to one deployment, reducing effective quota by the number of users +- **`session_affinity`**: Requires explicit session IDs and still reduces quota + +```mermaid +flowchart TD + A["1. Initial request to Responses API + router.aresponses()"] --> B["2. Router load balances to Deployment A + (API Key 1, Azure East US)"] + B --> C["3. Response contains encrypted item + rs_abc123 (encrypted with Org 1 key)"] + C --> D["4. Follow-up request includes rs_abc123 in input"] + D --> E["5. Router load balances to Deployment B + (API Key 2, Azure West Europe)"] + E -->|"Different API key"| F["6. ❌ Deployment B cannot decrypt rs_abc123 + Error: invalid_encrypted_content"] + + D -.->|"With encrypted_content_affinity"| G["5b. Router detects rs_abc123 was created by Deployment A"] + G --> H["6b. ✅ Routes to Deployment A (bypasses rate limits) + Request succeeds"] + + style F fill:#f8d7da,stroke:#dc3545 + style H fill:#d4edda,stroke:#28a745 + style E fill:#fff3cd,stroke:#ffc107 + style G fill:#d4edda,stroke:#28a745 +``` + +--- + +## Root Cause + +LiteLLM's router had no mechanism to track which deployment created specific encrypted content items and route follow-up requests accordingly. The router treated all deployments as interchangeable, leading to decryption failures when encrypted content crossed organizational boundaries. + +**The Problem Flow:** + +1. User calls `router.aresponses()` with model `gpt-5.1-codex` +2. Router load balances to Deployment A (Azure East US, API Key 1) +3. Response contains encrypted reasoning item `rs_abc123` (encrypted with Org 1's key) +4. User makes follow-up request with `rs_abc123` in the input +5. Router load balances to Deployment B (Azure West Europe, API Key 2) +6. Deployment B tries to decrypt `rs_abc123` with Org 2's key → **fails** + +**Why Existing Solutions Didn't Work:** + +- **`previous_response_id`**: Not provided by all clients (e.g., Codex) +- **`deployment_affinity`**: Pins *all* user requests to one deployment → reduces quota to 1/N where N = number of deployments +- **`session_affinity`**: Requires explicit session management and still reduces quota + +**Timeline:** + +1. Users configured multi-region Responses API load balancing with different API keys +2. Initial requests succeeded, but follow-up requests with encrypted content failed intermittently +3. Error rate correlated with number of deployments (more deployments = higher chance of routing to wrong one) +4. Investigation revealed encrypted content was organization-bound +5. Existing affinity mechanisms deemed unsuitable (quota reduction, missing `previous_response_id`) +6. New solution designed and implemented: `encrypted_content_affinity` + +--- + +## The Fix + +Implemented a new `encrypted_content_affinity` pre-call check that intelligently tracks encrypted content and routes follow-up requests **only when necessary**. + +### Implementation + +**1. New `EncryptedContentAffinityCheck` Class** ([`encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py)) + +```python +class EncryptedContentAffinityCheck(CustomLogger): + """ + Routes follow-up Responses API requests to the deployment that produced + the encrypted output items they reference. + """ + + async def async_log_success_event(self, kwargs, response_obj, ...): + """Track: Extract item IDs from response output, cache item_id → deployment_id""" + output = self._get_output_from_response(response_obj) + item_ids = self._extract_item_ids_from_output(output) + model_id = self._get_model_id_from_kwargs(kwargs) + + for item_id in item_ids: + await self.cache.async_set_cache( + f"encrypted_content_affinity:v1:{item_id}", + model_id, + ttl=86400, # 24 hours + ) + + async def async_filter_deployments(self, model, healthy_deployments, ...): + """Route: Check if input contains tracked items, pin to originating deployment""" + input_item_ids = self._extract_item_ids_from_input(request_kwargs.get("input")) + + for item_id in input_item_ids: + cached_model_id = await self.cache.async_get_cache(f"...:{item_id}") + if cached_model_id: + deployment = self._find_deployment_by_model_id( + healthy_deployments, cached_model_id + ) + if deployment: + # Signal to bypass rate limits (encrypted content must go here) + request_kwargs["_encrypted_content_affinity_pinned"] = True + return [deployment] + + return healthy_deployments # Normal load balancing +``` + +**2. Rate Limit Bypass** ([`router.py#L8656-L8660`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py#L8656-L8660)) + +When encrypted content requires a specific deployment, RPM/TPM limits are bypassed (the request would fail on any other deployment anyway): + +```python +# In async_get_available_deployment, after filtering healthy deployments: +if ( + request_kwargs.get("_encrypted_content_affinity_pinned") + and len(healthy_deployments) == 1 +): + return healthy_deployments[0] # Bypass routing strategy (RPM/TPM checks) +``` + +**3. Configuration** + +```yaml +router_settings: + routing_strategy: usage-based-routing-v2 + enable_pre_call_checks: true + optional_pre_call_checks: + - encrypted_content_affinity + deployment_affinity_ttl_seconds: 86400 # 24 hours +``` + +### Key Benefits + +✅ **No quota reduction**: Only pins requests containing tracked encrypted items +✅ **Bypasses rate limits**: When encrypted content requires a specific deployment, RPM/TPM limits don't block it +✅ **No `previous_response_id` required**: Works by tracking item IDs in response output +✅ **Globally safe**: Can be enabled for all models; non-Responses-API calls are unaffected +✅ **Surgical precision**: Normal requests continue to load balance freely + +--- + +## Remediation + +| # | Action | Status | Code | +|---|---|---|---| +| 1 | Create `EncryptedContentAffinityCheck` class with tracking and routing logic | ✅ Done | [`encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py) | +| 2 | Add `encrypted_content_affinity` to `OptionalPreCallChecks` type | ✅ Done | [`router.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/types/router.py) | +| 3 | Wire up check in `Router.add_optional_pre_call_checks` | ✅ Done | [`router.py#L8656-L8660`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py#L8656-L8660) | +| 4 | Implement rate limit bypass for affinity-pinned requests | ✅ Done | [`router.py#L8656-L8660`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py#L8656-L8660) | +| 5 | Unit tests: tracking, routing, no-op for non-Responses-API, RPM bypass | ✅ Done | [`test_encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py) | +| 6 | Documentation: Responses API guide, load balancing guide, config reference | ✅ Done | [Docs](https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing) | + +--- + +## Migration Guide + +### Before (Using `deployment_affinity`) + +```yaml +router_settings: + optional_pre_call_checks: + - deployment_affinity # ❌ Reduces quota by number of users +``` + +**Problem:** All requests from a user pin to one deployment, reducing effective quota to 1/N. + +### After (Using `encrypted_content_affinity`) + +```yaml +router_settings: + optional_pre_call_checks: + - encrypted_content_affinity # ✅ Only pins requests with encrypted content +``` + +**Benefit:** Normal requests load balance freely, only encrypted content requests pin when necessary. + +--- From a88a17796b385ad4321a39e75523613095b5f078 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 22:56:57 +0530 Subject: [PATCH 045/147] Fix logging and encrypted content extraction --- .../encrypted_content_affinity_check.py | 24 ++++++++++++------- .../test_encrypted_content_affinity_check.py | 6 +++-- 2 files changed, 19 insertions(+), 11 deletions(-) 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 cdf69e4b24c..9cab59ea0e0 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 @@ -76,26 +76,31 @@ class EncryptedContentAffinityCheck(CustomLogger): def _extract_item_ids_from_output( output: list, ) -> List[str]: - """Extract all item IDs from a Responses API output list.""" + """Extract item IDs from output items that contain encrypted_content.""" item_ids: List[str] = [] for item in output: item_id: Optional[str] = None + has_encrypted_content = False + if isinstance(item, dict): item_id = item.get("id") + has_encrypted_content = "encrypted_content" in item else: item_id = getattr(item, "id", None) - if item_id and isinstance(item_id, str): + has_encrypted_content = hasattr(item, "encrypted_content") + + if item_id and isinstance(item_id, str) and has_encrypted_content: item_ids.append(item_id) return item_ids @staticmethod def _extract_item_ids_from_input(request_input: Any) -> List[str]: """ - Extract item IDs from the ``input`` field of a Responses API request. + Extract item IDs from input items that contain encrypted_content. ``input`` can be: - a plain string -> no item IDs - - a list of items -> each item may have an ``id`` field + - a list of items -> only extract IDs from items with encrypted_content """ if not isinstance(request_input, list): return [] @@ -104,7 +109,8 @@ class EncryptedContentAffinityCheck(CustomLogger): for item in request_input: if isinstance(item, dict): item_id = item.get("id") - if item_id and isinstance(item_id, str): + has_encrypted_content = "encrypted_content" in item + if item_id and isinstance(item_id, str) and has_encrypted_content: item_ids.append(item_id) return item_ids @@ -191,13 +197,13 @@ class EncryptedContentAffinityCheck(CustomLogger): ttl=self.ttl_seconds, ) except Exception as e: - verbose_router_logger.debug( + verbose_router_logger.error( "EncryptedContentAffinityCheck: failed to cache item_id=%s error=%s", item_id, e, ) - verbose_router_logger.info( + verbose_router_logger.debug( "EncryptedContentAffinityCheck: cached %d item IDs -> deployment=%s", len(item_ids), model_id, @@ -247,7 +253,7 @@ class EncryptedContentAffinityCheck(CustomLogger): model_id=cached_model_id, ) if deployment is not None: - verbose_router_logger.info( + verbose_router_logger.debug( "EncryptedContentAffinityCheck: item_id=%s pinning -> deployment=%s", item_id, cached_model_id, @@ -257,7 +263,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ] = True return [deployment] - verbose_router_logger.info( + verbose_router_logger.debug( "EncryptedContentAffinityCheck: cached deployment=%s for item_id=%s " "not found in healthy_deployments", cached_model_id, 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 c70290b674d..17cd2162d99 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 @@ -48,6 +48,7 @@ async def test_encrypted_content_affinity_tracks_and_routes(): "type": "reasoning", "id": "rs_encrypted_item_456", "status": "completed", + "encrypted_content": "gAAAAABpnW_yEYmSNEyOG...", }, ], "parallel_tool_calls": True, @@ -118,7 +119,7 @@ async def test_encrypted_content_affinity_tracks_and_routes(): model=model_group, input=[ {"type": "message", "id": "msg_abc123", "role": "assistant"}, - {"type": "reasoning", "id": "rs_encrypted_item_456"}, + {"type": "reasoning", "id": "rs_encrypted_item_456", "encrypted_content": "gAAAAABpnW_yEYmSNEyOG..."}, ], ) second_model_id = second_response._hidden_params["model_id"] @@ -204,6 +205,7 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): "type": "reasoning", "id": "rs_encrypted_must_pin", "status": "completed", + "encrypted_content": "gAAAAABpnW_yEYmSNEyOG...", }, ], "usage": {"input_tokens": 5, "output_tokens": 10, "total_tokens": 15}, @@ -255,7 +257,7 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): second_response = await router.aresponses( model="openai.gpt-5.1-codex", input=[ - {"type": "reasoning", "id": "rs_encrypted_must_pin"}, + {"type": "reasoning", "id": "rs_encrypted_must_pin", "encrypted_content": "gAAAAABpnW_yEYmSNEyOG..."}, ], ) second_model_id = second_response._hidden_params["model_id"] From 18bf3f2df649c53486bf4e15c353e8464f112afa Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 23:06:54 +0530 Subject: [PATCH 046/147] Fix mock github test --- .../test_encrypted_content_affinity_check.py | 52 ++++++------------- 1 file changed, 17 insertions(+), 35 deletions(-) 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 17cd2162d99..cfebaab346f 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 @@ -4,6 +4,8 @@ import sys from unittest.mock import AsyncMock, patch import pytest +import respx +from httpx import Response sys.path.insert(0, os.path.abspath("../..")) @@ -143,7 +145,8 @@ async def test_encrypted_content_affinity_no_effect_on_chat_completions(): "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "gpt-3.5-turbo", - "api_key": os.environ.get("OPENAI_API_KEY", "test-key"), + "api_key": "test-key", + "mock_response": "Hello from chat completion!", }, "model_info": {"id": "chat-deployment-1"}, }, @@ -151,41 +154,20 @@ async def test_encrypted_content_affinity_no_effect_on_chat_completions(): optional_pre_call_checks=["encrypted_content_affinity"], ) - mock_chat_response = { - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": "gpt-3.5-turbo", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Hello!"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}, - } + # Multiple chat completion requests should work normally + response1 = await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello"}], + ) + response2 = await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello again"}], + ) - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post: - mock_post.return_value = MockResponse(mock_chat_response, 200) - - # Multiple chat completion requests should work normally - response1 = await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello"}], - ) - response2 = await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello again"}], - ) - - # Both should succeed (no affinity interference) - # Check that responses have IDs (litellm may modify them) - assert response1.id is not None - assert response2.id is not None + # Both should succeed (no affinity interference) + # Check that responses have IDs + assert response1.id is not None + assert response2.id is not None @pytest.mark.asyncio From adec115db82505f34c9c018238fc7cedc4e202b1 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 23:58:50 +0530 Subject: [PATCH 047/147] Fix logging for error --- .../pre_call_checks/encrypted_content_affinity_check.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 9cab59ea0e0..d632dc6e0b0 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 @@ -263,7 +263,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ] = True return [deployment] - verbose_router_logger.debug( + verbose_router_logger.error( "EncryptedContentAffinityCheck: cached deployment=%s for item_id=%s " "not found in healthy_deployments", cached_model_id, From 122f534d8762ff76e72dad0ff264bc0f51b6f7b5 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 27 Feb 2026 16:00:43 +0530 Subject: [PATCH 048/147] Add encoding method for Encrypted-content-aware deployment --- .../encrypted_content_affinity_check.py | 251 +++++------------- 1 file changed, 64 insertions(+), 187 deletions(-) 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 d632dc6e0b0..e6d691896ca 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 @@ -7,26 +7,36 @@ organization's API key. If a follow-up request containing those items is routed different deployment (different org), OpenAI rejects it with an `invalid_encrypted_content` error because the organization_id doesn't match. -This callback solves the problem by: -1. Tracking output item IDs from Responses API responses and mapping them to the - deployment (model_id) that produced them. -2. On subsequent requests, scanning the `input` field for known item IDs and pinning - the request to the originating deployment. +This callback solves the problem by encoding the originating deployment's ``model_id`` +directly into the item IDs of output items that carry ``encrypted_content`` (the same +approach used by the responses-API affinity for ``previous_response_id``). The encoded +ID is decoded on the next request so the router can pin to the correct deployment without +any cache lookup. + +Response post-processing (encoding) is handled by +``ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response`` which is +called inside ``_update_responses_api_response_id_with_model_id`` in ``responses/utils.py``. + +Request pre-processing (ID restoration before forwarding to upstream) is handled by +``ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input`` which is called +in ``get_optional_params_responses_api``. + +This pre-call check is responsible only for the routing decision: it reads the encoded +``model_id`` out of the item IDs and pins the request to the matching deployment. Safe to enable globally: -- Only activates when known item IDs appear in the request `input`. +- Only activates when encoded item IDs appear in the request ``input``. - No effect on embedding models, chat completions, or first-time requests. - No quota reduction -- first requests are fully load balanced. +- No cache required. """ from typing import Any, List, Optional, cast from litellm._logging import verbose_router_logger -from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse - -_DEFAULT_TTL_SECONDS = 86400 # 24 hours +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import AllMessageValues class EncryptedContentAffinityCheck(CustomLogger): @@ -34,89 +44,43 @@ class EncryptedContentAffinityCheck(CustomLogger): Routes follow-up Responses API requests to the deployment that produced the encrypted output items they reference. + 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"])``. """ - CACHE_KEY_PREFIX = "encrypted_content_affinity:v1" - - def __init__( - self, - cache: DualCache, - ttl_seconds: int = _DEFAULT_TTL_SECONDS, - ): + def __init__(self) -> None: super().__init__() - self.cache = cache - self.ttl_seconds = ttl_seconds # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ @staticmethod - def _get_output_from_response( - response_obj: Any, - ) -> Optional[list]: + def _extract_model_id_from_input(request_input: Any) -> Optional[str]: """ - Extract the ``output`` list from a Responses API response, handling - both ``ResponsesAPIResponse`` objects and plain dicts. - """ - if isinstance(response_obj, ResponsesAPIResponse): - return response_obj.output - if isinstance(response_obj, dict) and "output" in response_obj: - output = response_obj["output"] - if isinstance(output, list): - return output - if hasattr(response_obj, "output"): - output = response_obj.output - if isinstance(output, list): - return output - return None - - @staticmethod - def _extract_item_ids_from_output( - output: list, - ) -> List[str]: - """Extract item IDs from output items that contain encrypted_content.""" - item_ids: List[str] = [] - for item in output: - item_id: Optional[str] = None - has_encrypted_content = False - - if isinstance(item, dict): - item_id = item.get("id") - has_encrypted_content = "encrypted_content" in item - else: - item_id = getattr(item, "id", None) - has_encrypted_content = hasattr(item, "encrypted_content") - - if item_id and isinstance(item_id, str) and has_encrypted_content: - item_ids.append(item_id) - return item_ids - - @staticmethod - def _extract_item_ids_from_input(request_input: Any) -> List[str]: - """ - Extract item IDs from input items that contain encrypted_content. + Scan ``input`` items for litellm-encoded encrypted-content item IDs and + return the ``model_id`` embedded in the first one found. ``input`` can be: - - a plain string -> no item IDs - - a list of items -> only extract IDs from items with encrypted_content + - a plain string -> no encoded IDs + - a list of items -> check each item's ``id`` field """ if not isinstance(request_input, list): - return [] + return None - item_ids: List[str] = [] for item in request_input: - if isinstance(item, dict): - item_id = item.get("id") - has_encrypted_content = "encrypted_content" in item - if item_id and isinstance(item_id, str) and has_encrypted_content: - item_ids.append(item_id) - return item_ids + if not isinstance(item, dict): + continue + item_id = item.get("id") + if not item_id or not isinstance(item_id, str): + continue + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) + if decoded: + return decoded.get("model_id") - @classmethod - def _cache_key(cls, item_id: str) -> str: - return f"{cls.CACHE_KEY_PREFIX}:{item_id}" + return None @staticmethod def _find_deployment_by_model_id( @@ -133,82 +97,6 @@ class EncryptedContentAffinityCheck(CustomLogger): return deployment return None - @staticmethod - def _get_model_id_from_kwargs(kwargs: dict) -> Optional[str]: - """ - Extract the deployment model_id from success-callback kwargs. - - The Router populates ``litellm_params.metadata.model_info.id`` after - selecting a deployment. Also check top-level ``model_info`` as a - fallback (some call paths set it there). - """ - # Primary path: litellm_params -> metadata -> model_info -> id - litellm_params = kwargs.get("litellm_params") - if isinstance(litellm_params, dict): - metadata = litellm_params.get("metadata") - if isinstance(metadata, dict): - model_info = metadata.get("model_info") - if isinstance(model_info, dict): - model_id = model_info.get("id") - if model_id is not None: - return str(model_id) - - # Fallback: top-level model_info (set by some router call paths) - model_info = kwargs.get("model_info") - if isinstance(model_info, dict): - model_id = model_info.get("id") - if model_id is not None: - return str(model_id) - - return None - - # ------------------------------------------------------------------ - # Response tracking (success callback) - # ------------------------------------------------------------------ - - async def async_log_success_event( - self, kwargs: dict, response_obj: Any, start_time: Any, end_time: Any - ) -> None: - """ - After a successful Responses API call, cache each output item ID - mapped to the deployment that produced it. - """ - output = self._get_output_from_response(response_obj) - if output is None: - return - - model_id = self._get_model_id_from_kwargs(kwargs) - if not model_id: - verbose_router_logger.debug( - "EncryptedContentAffinityCheck: model_id not found in kwargs, skipping tracking", - ) - return - - item_ids = self._extract_item_ids_from_output(output) - if not item_ids: - return - - for item_id in item_ids: - try: - cache_key = self._cache_key(item_id) - await self.cache.async_set_cache( - cache_key, - model_id, - ttl=self.ttl_seconds, - ) - except Exception as e: - verbose_router_logger.error( - "EncryptedContentAffinityCheck: failed to cache item_id=%s error=%s", - item_id, - e, - ) - - verbose_router_logger.debug( - "EncryptedContentAffinityCheck: cached %d item IDs -> deployment=%s", - len(item_ids), - model_id, - ) - # ------------------------------------------------------------------ # Request routing (pre-call filter) # ------------------------------------------------------------------ @@ -222,52 +110,41 @@ class EncryptedContentAffinityCheck(CustomLogger): parent_otel_span: Optional[Span] = None, ) -> List[dict]: """ - If the request ``input`` contains items whose IDs were previously - tracked, pin the request to the deployment that produced them. + If the request ``input`` contains litellm-encoded item IDs, decode the + embedded ``model_id`` and pin the request to that deployment. """ request_kwargs = request_kwargs or {} typed_healthy_deployments = cast(List[dict], healthy_deployments) + # Signal to the response post-processor that encrypted item IDs should be + # encoded in the output of this request. + litellm_metadata = request_kwargs.setdefault("litellm_metadata", {}) + litellm_metadata["encrypted_content_affinity_enabled"] = True + request_input = request_kwargs.get("input") - input_item_ids = self._extract_item_ids_from_input(request_input) - if not input_item_ids: + model_id = self._extract_model_id_from_input(request_input) + if not model_id: return typed_healthy_deployments verbose_router_logger.debug( - "EncryptedContentAffinityCheck: found %d item IDs in input, checking cache", - len(input_item_ids), + "EncryptedContentAffinityCheck: decoded model_id=%s from input item IDs", + model_id, ) - for item_id in input_item_ids: - cache_key = self._cache_key(item_id) - try: - cached_model_id = await self.cache.async_get_cache(key=cache_key) - except Exception: - continue - - if not cached_model_id or not isinstance(cached_model_id, str): - continue - - deployment = self._find_deployment_by_model_id( - healthy_deployments=typed_healthy_deployments, - model_id=cached_model_id, - ) - if deployment is not None: - verbose_router_logger.debug( - "EncryptedContentAffinityCheck: item_id=%s pinning -> deployment=%s", - item_id, - cached_model_id, - ) - request_kwargs[ - "_encrypted_content_affinity_pinned" - ] = True - return [deployment] - - verbose_router_logger.error( - "EncryptedContentAffinityCheck: cached deployment=%s for item_id=%s " - "not found in healthy_deployments", - cached_model_id, - item_id, + deployment = self._find_deployment_by_model_id( + healthy_deployments=typed_healthy_deployments, + model_id=model_id, + ) + if deployment is not None: + verbose_router_logger.debug( + "EncryptedContentAffinityCheck: pinning -> deployment=%s", + model_id, ) + request_kwargs["_encrypted_content_affinity_pinned"] = True + return [deployment] + verbose_router_logger.error( + "EncryptedContentAffinityCheck: decoded deployment=%s not found in healthy_deployments", + model_id, + ) return typed_healthy_deployments From 7928d41e9ad93e8df8e1228a5bee301ab32258e4 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 27 Feb 2026 16:01:02 +0530 Subject: [PATCH 049/147] Update the routing --- litellm/responses/main.py | 11 ++++ litellm/responses/utils.py | 110 +++++++++++++++++++++++++++++++++++++ litellm/router.py | 5 +- 3 files changed, 122 insertions(+), 4 deletions(-) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 05fd6026af2..2576ed7db31 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -745,6 +745,11 @@ def responses( custom_llm_provider=custom_llm_provider, ) + # Decode any litellm-encoded encrypted-content item IDs back to their original IDs + input = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( + input + ) + # Call the handler with _is_async flag instead of directly calling the async handler response = base_llm_http_handler.response_api_handler( model=model, @@ -1617,6 +1622,12 @@ def compact_responses( custom_llm_provider=custom_llm_provider, ) + # Decode any litellm-encoded encrypted-content item IDs back to their original IDs + # before forwarding to the upstream provider. + input = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( + input + ) + # Call the handler with _is_async flag instead of directly calling the async handler response = base_llm_http_handler.compact_response_api_handler( model=model, diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 39aebb262fe..0c203dc6305 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -217,8 +217,118 @@ class ResponsesAPIRequestUtils: responses_api_response["id"] = updated_id else: responses_api_response.id = updated_id + + if litellm_metadata.get("encrypted_content_affinity_enabled"): + responses_api_response = ( + ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( + response=responses_api_response, + model_id=model_id, + ) + ) + return responses_api_response + @staticmethod + def _build_encrypted_item_id(model_id: str, item_id: str) -> str: + """Encode model_id into an output item ID for encrypted-content items. + + Format: ``encitem_{base64("litellm:model_id:{model_id};item_id:{original_id}")}`` + """ + assembled = f"litellm:model_id:{model_id};item_id:{item_id}" + encoded = base64.b64encode(assembled.encode("utf-8")).decode("utf-8") + return f"encitem_{encoded}" + + @staticmethod + def _decode_encrypted_item_id(encoded_id: str) -> Optional[Dict[str, str]]: + """Decode a litellm-encoded encrypted-content item ID. + + Returns a dict with ``model_id`` and ``item_id`` keys, or ``None`` if + the string is not a litellm-encoded item ID. + """ + if not encoded_id.startswith("encitem_"): + return None + try: + cleaned = encoded_id[len("encitem_"):] + # Restore any padding that may have been stripped in transit + missing = len(cleaned) % 4 + if missing: + cleaned += "=" * (4 - missing) + decoded = base64.b64decode(cleaned.encode("utf-8")).decode("utf-8") + # Split on first ";" only so that semicolons inside item_id are preserved + parts = decoded.split(";", 1) + if len(parts) < 2: + return None + model_id = parts[0].replace("litellm:model_id:", "") + item_id = parts[1].replace("item_id:", "") + return {"model_id": model_id, "item_id": item_id} + except Exception: + return None + + @staticmethod + def _update_encrypted_content_item_ids_in_response( + response: Union["ResponsesAPIResponse", Dict[str, Any]], + model_id: Optional[str], + ) -> Union["ResponsesAPIResponse", Dict[str, Any]]: + """Rewrite item IDs for output items that contain ``encrypted_content``. + + Encodes ``model_id`` into the item ID so that follow-up requests can be + routed back to the originating deployment without any cache lookup. + """ + if not model_id: + return response + + output: Optional[list] = None + if isinstance(response, dict): + output = response.get("output") + else: + output = getattr(response, "output", None) + + if not isinstance(output, list): + return response + + for item in output: + if isinstance(item, dict): + item_id = item.get("id") + if item_id and isinstance(item_id, str) and "encrypted_content" in item: + item["id"] = ResponsesAPIRequestUtils._build_encrypted_item_id( + model_id, item_id + ) + else: + item_id = getattr(item, "id", None) + if ( + item_id + and isinstance(item_id, str) + and hasattr(item, "encrypted_content") + ): + try: + item.id = ResponsesAPIRequestUtils._build_encrypted_item_id( + model_id, item_id + ) + except AttributeError: + pass + + return response + + @staticmethod + def _restore_encrypted_content_item_ids_in_input(request_input: Any) -> Any: + """Decode litellm-encoded item IDs in request input back to original IDs. + + Called before forwarding the request to the upstream provider so the + provider receives the original item IDs it issued. + """ + if not isinstance(request_input, list): + return request_input + + for item in request_input: + if isinstance(item, dict): + item_id = item.get("id") + if item_id and isinstance(item_id, str): + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) + if decoded: + item["id"] = decoded["item_id"] + + return request_input + @staticmethod def _build_responses_api_response_id( custom_llm_provider: Optional[str], diff --git a/litellm/router.py b/litellm/router.py index 7dee4fa83de..652cd68b555 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1263,10 +1263,7 @@ class Router: for cb in self.optional_callbacks ) if not already_registered: - ec_callback = EncryptedContentAffinityCheck( - cache=self.cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - ) + ec_callback = EncryptedContentAffinityCheck() self.optional_callbacks.append(ec_callback) litellm.logging_callback_manager.add_litellm_callback(ec_callback) From 37834f1d2a2adc5c80987f940f935fa56bbeb838 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 27 Feb 2026 16:01:23 +0530 Subject: [PATCH 050/147] Update the docs --- .../index.md | 90 +++++++++++-------- docs/my-website/docs/response_api.md | 34 +++---- 2 files changed, 62 insertions(+), 62 deletions(-) diff --git a/docs/my-website/blog/responses_api_encrypted_content_incident/index.md b/docs/my-website/blog/responses_api_encrypted_content_incident/index.md index 8422f6d3cd0..f229f9567da 100644 --- a/docs/my-website/blog/responses_api_encrypted_content_incident/index.md +++ b/docs/my-website/blog/responses_api_encrypted_content_incident/index.md @@ -119,47 +119,59 @@ Implemented a new `encrypted_content_affinity` pre-call check that intelligently ### Implementation -**1. New `EncryptedContentAffinityCheck` Class** ([`encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py)) +**1. Encoding `model_id` into output item IDs** ([`responses/utils.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/responses/utils.py)) + +The same approach used for `previous_response_id` affinity — no cache needed. When a response contains output items with `encrypted_content`, LiteLLM rewrites their IDs to embed the originating deployment's `model_id`: + +```python +# On response: rs_abc123 → encitem_{base64("litellm:model_id:{model_id};item_id:rs_abc123")} +def _build_encrypted_item_id(model_id: str, item_id: str) -> str: + assembled = f"litellm:model_id:{model_id};item_id:{item_id}" + encoded = base64.b64encode(assembled.encode("utf-8")).decode("utf-8") + return f"encitem_{encoded}" + +# On request: decode encitem_... → extract model_id for routing +def _decode_encrypted_item_id(encoded_id: str) -> Optional[Dict[str, str]]: + if not encoded_id.startswith("encitem_"): + return None + cleaned = encoded_id[len("encitem_"):] + missing = len(cleaned) % 4 + if missing: + cleaned += "=" * (4 - missing) # restore padding stripped in transit + decoded = base64.b64decode(cleaned).decode("utf-8") + model_id, item_id = decoded.split(";", 1) + return {"model_id": model_id.replace("litellm:model_id:", ""), + "item_id": item_id.replace("item_id:", "")} +``` + +Before forwarding to the upstream provider, LiteLLM restores the original item IDs so the provider never sees the encoded form: + +```python +# In responses/main.py — before calling the handler +input = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(input) +``` + +**2. `EncryptedContentAffinityCheck` — routing only** ([`encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py)) + +No `async_log_success_event` or cache lookups — the `model_id` is decoded directly from the item ID: ```python class EncryptedContentAffinityCheck(CustomLogger): - """ - Routes follow-up Responses API requests to the deployment that produced - the encrypted output items they reference. - """ - - async def async_log_success_event(self, kwargs, response_obj, ...): - """Track: Extract item IDs from response output, cache item_id → deployment_id""" - output = self._get_output_from_response(response_obj) - item_ids = self._extract_item_ids_from_output(output) - model_id = self._get_model_id_from_kwargs(kwargs) - - for item_id in item_ids: - await self.cache.async_set_cache( - f"encrypted_content_affinity:v1:{item_id}", - model_id, - ttl=86400, # 24 hours - ) - async def async_filter_deployments(self, model, healthy_deployments, ...): - """Route: Check if input contains tracked items, pin to originating deployment""" - input_item_ids = self._extract_item_ids_from_input(request_kwargs.get("input")) - - for item_id in input_item_ids: - cached_model_id = await self.cache.async_get_cache(f"...:{item_id}") - if cached_model_id: + """Decode encitem_ IDs in input to extract model_id and pin to that deployment.""" + for item in request_kwargs.get("input", []): + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item.get("id", "")) + if decoded: deployment = self._find_deployment_by_model_id( - healthy_deployments, cached_model_id + healthy_deployments, decoded["model_id"] ) if deployment: - # Signal to bypass rate limits (encrypted content must go here) request_kwargs["_encrypted_content_affinity_pinned"] = True return [deployment] - - return healthy_deployments # Normal load balancing + return healthy_deployments ``` -**2. Rate Limit Bypass** ([`router.py#L8656-L8660`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py#L8656-L8660)) +**3. Rate Limit Bypass** ([`router.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py)) When encrypted content requires a specific deployment, RPM/TPM limits are bypassed (the request would fail on any other deployment anyway): @@ -185,9 +197,10 @@ router_settings: ### Key Benefits -✅ **No quota reduction**: Only pins requests containing tracked encrypted items +✅ **No quota reduction**: Only pins requests containing encrypted items ✅ **Bypasses rate limits**: When encrypted content requires a specific deployment, RPM/TPM limits don't block it -✅ **No `previous_response_id` required**: Works by tracking item IDs in response output +✅ **No `previous_response_id` required**: Works by encoding `model_id` directly into the item ID +✅ **No cache required**: `model_id` is decoded on-the-fly from the item ID — no Redis, no TTL ✅ **Globally safe**: Can be enabled for all models; non-Responses-API calls are unaffected ✅ **Surgical precision**: Normal requests continue to load balance freely @@ -197,12 +210,13 @@ router_settings: | # | Action | Status | Code | |---|---|---|---| -| 1 | Create `EncryptedContentAffinityCheck` class with tracking and routing logic | ✅ Done | [`encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py) | -| 2 | Add `encrypted_content_affinity` to `OptionalPreCallChecks` type | ✅ Done | [`router.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/types/router.py) | -| 3 | Wire up check in `Router.add_optional_pre_call_checks` | ✅ Done | [`router.py#L8656-L8660`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py#L8656-L8660) | -| 4 | Implement rate limit bypass for affinity-pinned requests | ✅ Done | [`router.py#L8656-L8660`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py#L8656-L8660) | -| 5 | Unit tests: tracking, routing, no-op for non-Responses-API, RPM bypass | ✅ Done | [`test_encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py) | -| 6 | Documentation: Responses API guide, load balancing guide, config reference | ✅ Done | [Docs](https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing) | +| 1 | Encode `model_id` into encrypted-content item IDs on response | ✅ Done | [`responses/utils.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/responses/utils.py) | +| 2 | Restore original item IDs before forwarding to upstream provider | ✅ Done | [`responses/main.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/responses/main.py) | +| 3 | `EncryptedContentAffinityCheck`: decode item IDs to route (no cache) | ✅ Done | [`encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py) | +| 4 | Add `encrypted_content_affinity` to `OptionalPreCallChecks` type | ✅ Done | [`types/router.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/types/router.py) | +| 5 | Implement rate limit bypass for affinity-pinned requests | ✅ Done | [`router.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py) | +| 6 | Unit tests: encoding/decoding utilities, routing, RPM bypass | ✅ Done | [`test_encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py) | +| 7 | Documentation: Responses API guide, load balancing guide, config reference | ✅ Done | [Docs](https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing) | --- diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index 85144a80245..a7cf61ef16a 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -1011,24 +1011,25 @@ This error occurs when: ### The Solution: `encrypted_content_affinity` -The `encrypted_content_affinity` pre-call check intelligently tracks encrypted content and routes follow-up requests to the originating deployment **only when necessary**. +The `encrypted_content_affinity` pre-call check routes follow-up requests containing encrypted items to the originating deployment **only when necessary** **Key Benefits:** -- ✅ **No quota reduction**: Unlike `deployment_affinity`, only pins requests that contain tracked encrypted items +- ✅ **No quota reduction**: Unlike `deployment_affinity`, only pins requests that contain encrypted items - ✅ **Bypasses rate limits**: When encrypted content requires a specific deployment, RPM/TPM limits are bypassed (the request would fail on any other deployment anyway) -- ✅ **No `previous_response_id` required**: Works by tracking item IDs in response output and matching them in request input +- ✅ **No `previous_response_id` required**: Works by encoding `model_id` directly into item IDs +- ✅ **No cache required**: `model_id` is decoded on-the-fly — no Redis dependency, no TTL to manage - ✅ **Globally safe**: Can be enabled for all models; non-Responses-API calls (chat, embeddings) are unaffected ### How It Works -1. **Tracking Phase** (after successful response): - - Extracts all item IDs from response `output` (e.g., `msg_abc`, `rs_xyz`) - - Caches mapping: `item_id` → `deployment_id` (default TTL: 24 hours) +1. **Encoding Phase** (on response): + - For each output item that contains `encrypted_content`, LiteLLM rewrites the item ID to embed the originating `model_id`: `rs_xyz` → `encitem_{base64("litellm:model_id:{model_id};item_id:rs_xyz")}` + - The original item ID is restored before forwarding the request to the upstream provider 2. **Routing Phase** (before request): - - Scans request `input` for item IDs - - If tracked item found → pins to originating deployment, bypasses rate limits - - If no tracked items → normal load balancing + - Scans request `input` for `encitem_` prefixed IDs + - If found → decodes `model_id`, pins to originating deployment, bypasses rate limits + - If no encoded items → normal load balancing ### Configuration @@ -1058,7 +1059,6 @@ router = Router( }, ], optional_pre_call_checks=["encrypted_content_affinity"], - deployment_affinity_ttl_seconds=86400, # 24 hours (default) ) # Initial request - routes to any deployment @@ -1104,7 +1104,6 @@ router_settings: enable_pre_call_checks: true optional_pre_call_checks: - encrypted_content_affinity - deployment_affinity_ttl_seconds: 86400 # Optional, default is 86400 (24 hours) ``` **Start proxy:** @@ -1124,19 +1123,6 @@ litellm --config config.yaml | `session_affinity` | Session-based applications | All requests with same `session_id` | ⚠️ Reduces quota by # of sessions | | `deployment_affinity` | Simple sticky sessions | All requests from same API key | ❌ Reduces quota by # of users | -### Multi-Instance Deployment (Redis) - -For multiple LiteLLM proxy instances, use Redis to share affinity state: - -```yaml -router_settings: - optional_pre_call_checks: - - encrypted_content_affinity - redis_host: redis.example.com - redis_port: 6379 - redis_password: your-password -``` - ## Calling non-Responses API endpoints (`/responses` to `/chat/completions` Bridge) From 2bc4da76ce62070bd753e9a3593c879060da389f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 27 Feb 2026 16:01:38 +0530 Subject: [PATCH 051/147] Update the tests --- .../test_encrypted_content_affinity_check.py | 242 ++++++++++++++---- 1 file changed, 195 insertions(+), 47 deletions(-) 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 cfebaab346f..66177208f4f 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 @@ -1,17 +1,31 @@ -import asyncio +""" +Tests for encrypted_content_affinity pre-call check. + +The mechanism works without any cache: +- On response: item IDs for output items with `encrypted_content` are rewritten to + `encitem_{base64("litellm:model_id:{model_id};item_id:{original_id}")}`. +- On routing: `EncryptedContentAffinityCheck` decodes the `encitem_` prefix to extract + `model_id` and pins the request to that deployment. +- Before forwarding: `_restore_encrypted_content_item_ids_in_input` decodes the IDs back + to their original form before sending to the upstream provider. +""" + import os import sys from unittest.mock import AsyncMock, patch import pytest -import respx -from httpx import Response sys.path.insert(0, os.path.abspath("../..")) import json import litellm +from litellm.responses.utils import ResponsesAPIRequestUtils + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- class MockResponse: @@ -25,12 +39,147 @@ class MockResponse: return self._json_data +def _get_item_id(item) -> str: + """Extract item ID from either a Pydantic model or a dict.""" + if isinstance(item, dict): + return item.get("id", "") + return getattr(item, "id", "") or "" + + +def _has_encrypted_content(item) -> bool: + """Check whether an output item carries encrypted_content.""" + if isinstance(item, dict): + return "encrypted_content" in item + return hasattr(item, "encrypted_content") and getattr(item, "encrypted_content") is not None + + +def _extract_encoded_item_id(response) -> str: + """ + Walk the response output and return the first litellm-encoded item ID + (i.e. one that starts with ``encitem_``). + """ + for item in response.output or []: + item_id = _get_item_id(item) + if item_id.startswith("encitem_"): + return item_id + return "" + + +# --------------------------------------------------------------------------- +# Unit tests for encoding / decoding utilities +# --------------------------------------------------------------------------- + + +class TestEncryptedItemIdCodec: + def test_roundtrip(self): + model_id = "deployment-1" + original_item_id = "rs_abc123def456" + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) + assert encoded.startswith("encitem_") + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded) + assert decoded is not None + assert decoded["model_id"] == model_id + assert decoded["item_id"] == original_item_id + + def test_decode_without_padding(self): + """Decoding must succeed even if base64 padding (=) was stripped in transit.""" + model_id = "gpt-5.1-codex-openai-2" + original_item_id = "rs_0efb96cb222403210069a01d5d52588196a9dc394ffdb89d00" + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) + # Strip any trailing '=' to simulate what happens in transit + stripped = encoded.rstrip("=") + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(stripped) + assert decoded is not None + assert decoded["model_id"] == model_id + assert decoded["item_id"] == original_item_id + + def test_non_encoded_id_returns_none(self): + assert ResponsesAPIRequestUtils._decode_encrypted_item_id("rs_abc123") is None + assert ResponsesAPIRequestUtils._decode_encrypted_item_id("msg_abc") is None + assert ResponsesAPIRequestUtils._decode_encrypted_item_id("") is None + + def test_semicolon_in_item_id(self): + """item_id values containing ';' must survive the roundtrip.""" + model_id = "deployment-1" + original_item_id = "rs_part1;part2;part3" + encoded = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_item_id) + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded) + assert decoded is not None + assert decoded["item_id"] == original_item_id + + +class TestUpdateEncryptedContentItemIds: + def test_rewrites_encrypted_items_in_dict_response(self): + model_id = "deployment-1" + response = { + "id": "resp_123", + "output": [ + {"id": "msg_abc", "type": "message", "content": []}, + {"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}, + ], + } + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( + response, model_id + ) + # Plain message item untouched + assert result["output"][0]["id"] == "msg_abc" + # Reasoning item with encrypted_content gets encoded + encoded_id = result["output"][1]["id"] + assert encoded_id.startswith("encitem_") + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded_id) + assert decoded["model_id"] == model_id + assert decoded["item_id"] == "rs_xyz" + + def test_no_op_when_model_id_is_none(self): + response = { + "output": [{"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}] + } + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( + response, None + ) + assert result["output"][0]["id"] == "rs_xyz" + + +class TestRestoreEncryptedContentItemIds: + def test_restores_encoded_ids(self): + model_id = "deployment-1" + original_id = "rs_encrypted_item_456" + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, original_id) + + request_input = [ + {"type": "message", "id": "msg_abc123", "role": "assistant"}, + {"type": "reasoning", "id": encoded_id, "encrypted_content": "secret"}, + ] + restored = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( + request_input + ) + assert restored[0]["id"] == "msg_abc123" + assert restored[1]["id"] == original_id + + def test_no_op_for_plain_string_input(self): + result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( + "Hello world" + ) + assert result == "Hello world" + + def test_no_op_for_unencoded_ids(self): + request_input = [{"type": "message", "id": "msg_plain"}] + result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( + request_input + ) + assert result[0]["id"] == "msg_plain" + + +# --------------------------------------------------------------------------- +# Integration tests (router-level) +# --------------------------------------------------------------------------- + + @pytest.mark.asyncio async def test_encrypted_content_affinity_tracks_and_routes(): """ - When encrypted_content_affinity is enabled, output item IDs from responses - are tracked, and follow-up requests containing those IDs route to the same - deployment. + The first response rewrites encrypted-content item IDs to encoded form. + The follow-up request with those encoded IDs is pinned to the same deployment. """ mock_response_data = { "id": "resp_mock-123", @@ -54,11 +203,7 @@ async def test_encrypted_content_affinity_tracks_and_routes(): }, ], "parallel_tool_calls": True, - "usage": { - "input_tokens": 5, - "output_tokens": 10, - "total_tokens": 15, - }, + "usage": {"input_tokens": 5, "output_tokens": 10, "total_tokens": 15}, "error": None, } @@ -84,14 +229,9 @@ async def test_encrypted_content_affinity_tracks_and_routes(): optional_pre_call_checks=["encrypted_content_affinity"], ) - model_group = "openai.gpt-5.1-codex" - - # Track which deployment was selected selected_deployments = [] def deterministic_choice(seq): - # First call: select deployment-1 - # Second call: would select deployment-2, but affinity should override if len(selected_deployments) == 0: return seq[0] return seq[1] if len(seq) > 1 else seq[0] @@ -105,39 +245,49 @@ async def test_encrypted_content_affinity_tracks_and_routes(): ): mock_post.return_value = MockResponse(mock_response_data, 200) - # First request: no encrypted items in input + # First request — goes to deployment-1 via deterministic_choice first_response = await router.aresponses( - model=model_group, + model="openai.gpt-5.1-codex", input="Hello, how are you?", ) first_model_id = first_response._hidden_params["model_id"] selected_deployments.append(first_model_id) - # Give async callbacks time to run - await asyncio.sleep(0.2) + # The response must have rewritten the encrypted item's ID to encoded form + encoded_item_id = _extract_encoded_item_id(first_response) + assert encoded_item_id.startswith("encitem_"), ( + f"Expected output item ID to be rewritten to encitem_... but got {encoded_item_id!r}" + ) - # Second request: includes encrypted item IDs from first response + # Verify the encoded ID decodes back to the correct deployment + original ID + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(encoded_item_id) + assert decoded is not None + assert decoded["model_id"] == first_model_id + assert decoded["item_id"] == "rs_encrypted_item_456" + + # Second request: use the encoded item IDs from the first response second_response = await router.aresponses( - model=model_group, + model="openai.gpt-5.1-codex", input=[ {"type": "message", "id": "msg_abc123", "role": "assistant"}, - {"type": "reasoning", "id": "rs_encrypted_item_456", "encrypted_content": "gAAAAABpnW_yEYmSNEyOG..."}, + { + "type": "reasoning", + "id": encoded_item_id, + "encrypted_content": "gAAAAABpnW_yEYmSNEyOG...", + }, ], ) second_model_id = second_response._hidden_params["model_id"] - # Affinity should route to the same deployment assert second_model_id == first_model_id, ( - f"Expected affinity to route to {first_model_id}, " - f"but got {second_model_id}" + f"Expected affinity to route to {first_model_id}, but got {second_model_id}" ) @pytest.mark.asyncio async def test_encrypted_content_affinity_no_effect_on_chat_completions(): """ - Encrypted content affinity should not affect regular chat completions - (they don't use the Responses API). + Encrypted content affinity should not affect regular chat completions. """ router = litellm.Router( model_list=[ @@ -154,7 +304,6 @@ async def test_encrypted_content_affinity_no_effect_on_chat_completions(): optional_pre_call_checks=["encrypted_content_affinity"], ) - # Multiple chat completion requests should work normally response1 = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello"}], @@ -163,9 +312,6 @@ async def test_encrypted_content_affinity_no_effect_on_chat_completions(): model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello again"}], ) - - # Both should succeed (no affinity interference) - # Check that responses have IDs assert response1.id is not None assert response2.id is not None @@ -173,8 +319,8 @@ async def test_encrypted_content_affinity_no_effect_on_chat_completions(): @pytest.mark.asyncio async def test_encrypted_content_affinity_bypasses_rpm_limits(): """ - When encrypted content affinity pins to a deployment, it should bypass - RPM limits since the encrypted content will fail on any other deployment. + When encrypted content affinity pins to a deployment, RPM limits are bypassed + since the request would fail on any other deployment anyway. """ mock_response_data = { "id": "resp_mock-rpm-test", @@ -225,34 +371,40 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): ) as mock_post: mock_post.return_value = MockResponse(mock_response_data, 200) - # First request goes to the low-RPM deployment first_response = await router.aresponses( model="openai.gpt-5.1-codex", input="Initial request", ) first_model_id = first_response._hidden_params["model_id"] - await asyncio.sleep(0.2) + # Extract encoded item ID from the first response output + encoded_item_id = _extract_encoded_item_id(first_response) + assert encoded_item_id.startswith("encitem_"), ( + f"Expected encitem_... but got {encoded_item_id!r}" + ) - # Second request with encrypted content should pin to the same deployment - # even though it's at RPM limit + # Follow-up with the encoded item ID — should pin to same deployment + # even if it is at its RPM limit second_response = await router.aresponses( model="openai.gpt-5.1-codex", input=[ - {"type": "reasoning", "id": "rs_encrypted_must_pin", "encrypted_content": "gAAAAABpnW_yEYmSNEyOG..."}, + { + "type": "reasoning", + "id": encoded_item_id, + "encrypted_content": "gAAAAABpnW_yEYmSNEyOG...", + }, ], ) second_model_id = second_response._hidden_params["model_id"] - # Should route to the same deployment despite RPM limit assert second_model_id == first_model_id @pytest.mark.asyncio async def test_encrypted_content_affinity_no_match_normal_routing(): """ - When input contains item IDs that aren't tracked, normal load balancing - should occur. + Input items with non-encoded IDs (no encitem_ prefix) fall through to + normal load balancing. """ mock_response_data = { "id": "resp_mock-no-match", @@ -301,15 +453,11 @@ async def test_encrypted_content_affinity_no_match_normal_routing(): ) as mock_post: mock_post.return_value = MockResponse(mock_response_data, 200) - # Request with unknown item IDs should use normal routing + # Non-encoded item ID — no affinity should kick in response = await router.aresponses( model="openai.gpt-5.1-codex", input=[ {"type": "message", "id": "unknown_item_id_12345"}, ], ) - - # Should succeed with normal routing (litellm may modify the ID) assert response.id is not None - # Verify it contains the original response ID in some form - assert "resp_mock-no-match" in str(response.id) or response.id.startswith("resp_") From 521f804350069acbbf88025d77e805b35d41affa Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 18:13:03 +0530 Subject: [PATCH 052/147] Fix encrypted content streaming affinity issue --- .../index.md | 118 ++++++-- .../exception_mapping_utils.py | 45 ++- litellm/responses/streaming_iterator.py | 30 +- litellm/responses/utils.py | 108 ++++++- .../encrypted_content_affinity_check.py | 52 +++- .../azure/test_azure_exception_mapping.py | 57 +++- .../test_encrypted_content_affinity_check.py | 280 +++++++++++++++++- 7 files changed, 624 insertions(+), 66 deletions(-) diff --git a/docs/my-website/blog/responses_api_encrypted_content_incident/index.md b/docs/my-website/blog/responses_api_encrypted_content_incident/index.md index f229f9567da..19b55898caa 100644 --- a/docs/my-website/blog/responses_api_encrypted_content_incident/index.md +++ b/docs/my-website/blog/responses_api_encrypted_content_incident/index.md @@ -119,32 +119,36 @@ Implemented a new `encrypted_content_affinity` pre-call check that intelligently ### Implementation -**1. Encoding `model_id` into output item IDs** ([`responses/utils.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/responses/utils.py)) +**1. Encoding `model_id` into output items** ([`responses/utils.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/responses/utils.py)) -The same approach used for `previous_response_id` affinity — no cache needed. When a response contains output items with `encrypted_content`, LiteLLM rewrites their IDs to embed the originating deployment's `model_id`: +The same approach used for `previous_response_id` affinity — no cache needed. When a response contains output items with `encrypted_content`, LiteLLM encodes the originating deployment's `model_id` in **two places** for redundancy: + +1. **Into the item ID** (if present): `rs_abc123` → `encitem_{base64("litellm:model_id:{model_id};item_id:rs_abc123")}` +2. **Into the encrypted_content itself**: Wraps the content with `litellm_enc:{base64("model_id:{model_id}")};{original_encrypted_content}` ```python -# On response: rs_abc123 → encitem_{base64("litellm:model_id:{model_id};item_id:rs_abc123")} +# Encoding item IDs (when present) def _build_encrypted_item_id(model_id: str, item_id: str) -> str: assembled = f"litellm:model_id:{model_id};item_id:{item_id}" encoded = base64.b64encode(assembled.encode("utf-8")).decode("utf-8") return f"encitem_{encoded}" -# On request: decode encitem_... → extract model_id for routing -def _decode_encrypted_item_id(encoded_id: str) -> Optional[Dict[str, str]]: - if not encoded_id.startswith("encitem_"): - return None - cleaned = encoded_id[len("encitem_"):] - missing = len(cleaned) % 4 - if missing: - cleaned += "=" * (4 - missing) # restore padding stripped in transit - decoded = base64.b64decode(cleaned).decode("utf-8") - model_id, item_id = decoded.split(";", 1) - return {"model_id": model_id.replace("litellm:model_id:", ""), - "item_id": item_id.replace("item_id:", "")} +# Wrapping encrypted_content (always, for redundancy) +def _wrap_encrypted_content_with_model_id(encrypted_content: str, model_id: str) -> str: + metadata = f"model_id:{model_id}" + encoded_metadata = base64.b64encode(metadata.encode("utf-8")).decode("utf-8") + return f"litellm_enc:{encoded_metadata};{encrypted_content}" ``` -Before forwarding to the upstream provider, LiteLLM restores the original item IDs so the provider never sees the encoded form: +**Why wrap encrypted_content directly?** Some clients (like Codex) don't consistently send item IDs in follow-up requests, but they always send the `encrypted_content` itself. By embedding `model_id` into the content, affinity works even when IDs are missing. + +**Streaming responses:** The wrapping logic is applied to both: +- Final response objects (non-streaming) +- Individual streaming events (`response.output_item.added`, `response.output_item.done`) + +This ensures clients receiving streaming responses get wrapped content they can send back. + +Before forwarding to the upstream provider, LiteLLM restores the original item IDs and unwraps encrypted_content so the provider never sees the encoded form: ```python # In responses/main.py — before calling the handler @@ -153,22 +157,43 @@ input = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(in **2. `EncryptedContentAffinityCheck` — routing only** ([`encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py)) -No `async_log_success_event` or cache lookups — the `model_id` is decoded directly from the item ID: +No `async_log_success_event` or cache lookups — the `model_id` is decoded directly from the item ID or encrypted_content: ```python class EncryptedContentAffinityCheck(CustomLogger): async def async_filter_deployments(self, model, healthy_deployments, ...): - """Decode encitem_ IDs in input to extract model_id and pin to that deployment.""" + """Extract model_id from input items (ID or encrypted_content) and pin to that deployment.""" for item in request_kwargs.get("input", []): - decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item.get("id", "")) - if decoded: + # Try to extract model_id from two sources: + model_id = self._extract_model_id_from_input(item) + + if model_id: deployment = self._find_deployment_by_model_id( - healthy_deployments, decoded["model_id"] + healthy_deployments, model_id ) if deployment: request_kwargs["_encrypted_content_affinity_pinned"] = True return [deployment] return healthy_deployments + + def _extract_model_id_from_input(self, item: dict) -> Optional[str]: + """Extract model_id from either encoded ID or wrapped encrypted_content.""" + # 1. Try decoding from item ID (if present) + item_id = item.get("id", "") + if item_id: + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) + if decoded: + return decoded["model_id"] + + # 2. Try unwrapping from encrypted_content (fallback for clients that omit IDs) + encrypted_content = item.get("encrypted_content", "") + if encrypted_content and encrypted_content.startswith("litellm_enc:"): + model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( + encrypted_content + ) + return model_id + + return None ``` **3. Rate Limit Bypass** ([`router.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py)) @@ -217,6 +242,57 @@ router_settings: | 5 | Implement rate limit bypass for affinity-pinned requests | ✅ Done | [`router.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/router.py) | | 6 | Unit tests: encoding/decoding utilities, routing, RPM bypass | ✅ Done | [`test_encrypted_content_affinity_check.py`](https://github.com/BerriAI/litellm/blob/main/litellm/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py) | | 7 | Documentation: Responses API guide, load balancing guide, config reference | ✅ Done | [Docs](https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing) | +| 8 | **[Mar 3]** Fix streaming events to wrap encrypted_content | ✅ Done | [`responses/streaming_iterator.py`](https://github.com/BerriAI/litellm/blob/main/litellm/litellm/responses/streaming_iterator.py) | + +--- + +## Follow-up Fix: Streaming Responses (Mar 3, 2026) + +### The Issue + +After the initial fix was deployed, users reported that the `invalid_encrypted_content` error **still occurred** when using streaming responses with clients like Codex. Investigation revealed: + +- ✅ Non-streaming responses: `encrypted_content` was correctly wrapped with `litellm_enc:` prefix +- ❌ Streaming responses: Individual `response.output_item.added` and `response.output_item.done` events contained **raw, unwrapped** `encrypted_content` + +Since Codex and other clients consume responses as streams, they received unwrapped content in these events and sent it back in follow-up requests, causing the affinity check to fail. + +### The Root Cause + +The `_update_encrypted_content_item_ids_in_response` function only modified the **final** response object, which is used for non-streaming responses. For streaming responses, individual chunks are processed by `ResponsesAPIStreamingIterator._process_chunk`, which was **not** applying the wrapping logic to streaming events. + +### The Fix + +Modified `litellm/litellm/responses/streaming_iterator.py` to wrap `encrypted_content` in streaming events: + +```python +# In ResponsesAPIStreamingIterator._process_chunk +if ( + self.litellm_metadata + and self.litellm_metadata.get("encrypted_content_affinity_enabled") +): + event_type = getattr(openai_responses_api_chunk, "type", None) + if event_type in ( + ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + ): + item = getattr(openai_responses_api_chunk, "item", None) + if item: + encrypted_content = getattr(item, "encrypted_content", None) + if encrypted_content and isinstance(encrypted_content, str): + model_id = ( + self.litellm_metadata.get("model_info", {}).get("id") + if self.litellm_metadata + else None + ) + if model_id: + wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id + ) + setattr(item, "encrypted_content", wrapped_content) +``` + +This ensures that **all** `encrypted_content` sent to clients (streaming or non-streaming) is wrapped with `model_id` metadata, enabling consistent affinity routing. --- diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index dde44cced36..951485130b3 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,9 +1,9 @@ import json +import re import traceback from typing import Any, Optional import httpx -import re import litellm from litellm._logging import verbose_logger @@ -443,6 +443,27 @@ def exception_type( # type: ignore # noqa: PLR0915 response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, ) + elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str: + exception_mapping_worked = True + helpful_message = ( + f"{exception_provider} - {message}\n\n" + " This error occurs when load balancing Responses API across deployments with different API keys.\n" + " Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n" + " Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n" + " router_settings:\n" + " enable_pre_call_checks: true\n" + " optional_pre_call_checks:\n" + " - encrypted_content_affinity\n\n" + " Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing" + ) + raise BadRequestError( + message=helpful_message, + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + ) elif ( "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str @@ -2126,7 +2147,27 @@ def exception_type( # type: ignore # noqa: PLR0915 extra_information=extra_information, original_exception=original_exception, ) - + elif azure_error_code == "invalid_encrypted_content" or "could not be verified" in error_str: + exception_mapping_worked = True + helpful_message = ( + f"AzureException - {message}\n\n" + "This error occurs when load balancing Responses API across deployments with different API keys.\n" + " Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n" + " Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n" + " router_settings:\n" + " enable_pre_call_checks: true\n" + " optional_pre_call_checks:\n" + " - encrypted_content_affinity\n\n" + " Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing" + ) + raise BadRequestError( + message=helpful_message, + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + body=getattr(original_exception, "body", None), + ) elif "invalid_request_error" in error_str: exception_mapping_worked = True raise BadRequestError( diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 43ef4610b4b..f61f108c992 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -8,7 +8,10 @@ from typing import Any, Dict, Optional import httpx import litellm -from litellm.constants import LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING +from litellm.constants import ( + LITELLM_MAX_STREAMING_DURATION_SECONDS, + STREAM_SSE_DONE_STRING, +) from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -137,6 +140,31 @@ class BaseResponsesAPIStreamingIterator: ) setattr(openai_responses_api_chunk, "response", response) + # Wrap encrypted_content in streaming events (output_item.added, output_item.done) + if ( + self.litellm_metadata + and self.litellm_metadata.get("encrypted_content_affinity_enabled") + ): + event_type = getattr(openai_responses_api_chunk, "type", None) + if event_type in ( + ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + ): + item = getattr(openai_responses_api_chunk, "item", None) + if item: + encrypted_content = getattr(item, "encrypted_content", None) + if encrypted_content and isinstance(encrypted_content, str): + model_id = ( + self.litellm_metadata.get("model_info", {}).get("id") + if self.litellm_metadata + else None + ) + if model_id: + wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id + ) + setattr(item, "encrypted_content", wrapped_content) + # Store the completed response if ( openai_responses_api_chunk diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 0c203dc6305..89e89711706 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -264,6 +264,56 @@ class ResponsesAPIRequestUtils: except Exception: return None + @staticmethod + def _wrap_encrypted_content_with_model_id( + encrypted_content: str, model_id: str + ) -> str: + """Wrap encrypted_content with model_id metadata for affinity routing. + + When Codex or other clients send items with encrypted_content but no ID, + we encode the model_id directly into the encrypted_content itself. + + Format: ``litellm_enc:{base64("model_id:{model_id}")};{original_encrypted_content}`` + """ + metadata = f"model_id:{model_id}" + encoded_metadata = base64.b64encode(metadata.encode("utf-8")).decode("utf-8") + return f"litellm_enc:{encoded_metadata};{encrypted_content}" + + @staticmethod + def _unwrap_encrypted_content_with_model_id( + wrapped_content: str, + ) -> tuple[Optional[str], str]: + """Unwrap encrypted_content to extract model_id and original content. + + Returns: + Tuple of (model_id, original_encrypted_content). + If not wrapped, returns (None, original_content). + """ + if not wrapped_content.startswith("litellm_enc:"): + return None, wrapped_content + + try: + # Split on first ";" to separate metadata from content + parts = wrapped_content.split(";", 1) + if len(parts) < 2: + return None, wrapped_content + + metadata_b64 = parts[0].replace("litellm_enc:", "") + original_content = parts[1] + + # Restore padding if needed + missing = len(metadata_b64) % 4 + if missing: + metadata_b64 += "=" * (4 - missing) + + decoded_metadata = base64.b64decode(metadata_b64.encode("utf-8")).decode( + "utf-8" + ) + model_id = decoded_metadata.replace("model_id:", "") + return model_id, original_content + except Exception: + return None, wrapped_content + @staticmethod def _update_encrypted_content_item_ids_in_response( response: Union["ResponsesAPIResponse", Dict[str, Any]], @@ -273,6 +323,9 @@ class ResponsesAPIRequestUtils: Encodes ``model_id`` into the item ID so that follow-up requests can be routed back to the originating deployment without any cache lookup. + + For items without an ID (e.g., from Codex), encodes model_id directly + into the encrypted_content itself. """ if not model_id: return response @@ -289,23 +342,42 @@ class ResponsesAPIRequestUtils: for item in output: if isinstance(item, dict): item_id = item.get("id") - if item_id and isinstance(item_id, str) and "encrypted_content" in item: - item["id"] = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, item_id + encrypted_content = item.get("encrypted_content") + + if encrypted_content and isinstance(encrypted_content, str): + # Always wrap encrypted_content with model_id for redundancy + item["encrypted_content"] = ( + ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id + ) ) + # Also encode the ID if present + if item_id and isinstance(item_id, str): + item["id"] = ResponsesAPIRequestUtils._build_encrypted_item_id( + model_id, item_id + ) else: item_id = getattr(item, "id", None) - if ( - item_id - and isinstance(item_id, str) - and hasattr(item, "encrypted_content") - ): + encrypted_content = getattr(item, "encrypted_content", None) + + if encrypted_content and isinstance(encrypted_content, str): + # Always wrap encrypted_content with model_id for redundancy try: - item.id = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, item_id + item.encrypted_content = ( + ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id + ) ) except AttributeError: pass + # Also encode the ID if present + if item_id and isinstance(item_id, str): + try: + item.id = ResponsesAPIRequestUtils._build_encrypted_item_id( + model_id, item_id + ) + except AttributeError: + pass return response @@ -314,7 +386,11 @@ class ResponsesAPIRequestUtils: """Decode litellm-encoded item IDs in request input back to original IDs. Called before forwarding the request to the upstream provider so the - provider receives the original item IDs it issued. + provider receives the original item IDs and unwrapped encrypted_content. + + Handles both: + 1. Items with encoded IDs (encitem_...) + 2. Items with wrapped encrypted_content (litellm_enc:...) """ if not isinstance(request_input, list): return request_input @@ -327,6 +403,16 @@ class ResponsesAPIRequestUtils: if decoded: item["id"] = decoded["item_id"] + encrypted_content = item.get("encrypted_content") + if encrypted_content and isinstance(encrypted_content, str): + _, unwrapped = ( + ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( + encrypted_content + ) + ) + if unwrapped != encrypted_content: + item["encrypted_content"] = unwrapped + return request_input @staticmethod 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 e6d691896ca..dc44ef13b7c 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 @@ -8,24 +8,29 @@ different deployment (different org), OpenAI rejects it with an `invalid_encrypt error because the organization_id doesn't match. This callback solves the problem by encoding the originating deployment's ``model_id`` -directly into the item IDs of output items that carry ``encrypted_content`` (the same -approach used by the responses-API affinity for ``previous_response_id``). The encoded -ID is decoded on the next request so the router can pin to the correct deployment without -any cache lookup. +into the response output items that carry ``encrypted_content``. Two encoding strategies: + +1. **Items with IDs**: Encode model_id into the item ID itself (e.g., ``encitem_...``) +2. **Items without IDs** (Codex): Wrap the encrypted_content with model_id metadata + (e.g., ``litellm_enc:{base64_metadata};{original_encrypted_content}``) + +The encoded model_id is decoded on the next request so the router can pin to the correct +deployment without any cache lookup. Response post-processing (encoding) is handled by ``ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response`` which is called inside ``_update_responses_api_response_id_with_model_id`` in ``responses/utils.py``. -Request pre-processing (ID restoration before forwarding to upstream) is handled by +Request pre-processing (ID/content restoration before forwarding to upstream) is handled by ``ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input`` which is called in ``get_optional_params_responses_api``. This pre-call check is responsible only for the routing decision: it reads the encoded -``model_id`` out of the item IDs and pins the request to the matching deployment. +``model_id`` from either item IDs or wrapped encrypted_content and pins the request to +the matching deployment. Safe to enable globally: -- Only activates when encoded item IDs appear in the request ``input``. +- Only activates when encoded markers appear in the request ``input``. - No effect on embedding models, chat completions, or first-time requests. - No quota reduction -- first requests are fully load balanced. - No cache required. @@ -60,12 +65,16 @@ class EncryptedContentAffinityCheck(CustomLogger): @staticmethod def _extract_model_id_from_input(request_input: Any) -> Optional[str]: """ - Scan ``input`` items for litellm-encoded encrypted-content item IDs and + Scan ``input`` items for litellm-encoded encrypted-content markers and return the ``model_id`` embedded in the first one found. + Checks both: + 1. Encoded item IDs (encitem_...) - for clients that send IDs + 2. Wrapped encrypted_content (litellm_enc:...) - for clients like Codex that don't send IDs + ``input`` can be: - - a plain string -> no encoded IDs - - a list of items -> check each item's ``id`` field + - a plain string -> no encoded markers + - a list of items -> check each item's ``id`` and ``encrypted_content`` fields """ if not isinstance(request_input, list): return None @@ -73,12 +82,25 @@ class EncryptedContentAffinityCheck(CustomLogger): for item in request_input: if not isinstance(item, dict): continue + + # First, try to decode from item ID (if present) item_id = item.get("id") - if not item_id or not isinstance(item_id, str): - continue - decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) - if decoded: - return decoded.get("model_id") + if item_id and isinstance(item_id, str): + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) + if decoded: + return decoded.get("model_id") + + # If no encoded ID, check if encrypted_content itself is wrapped + encrypted_content = item.get("encrypted_content") + if encrypted_content and isinstance(encrypted_content, str): + ( + model_id, + _, + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( + encrypted_content + ) + if model_id: + return model_id return None diff --git a/tests/test_litellm/llms/azure/test_azure_exception_mapping.py b/tests/test_litellm/llms/azure/test_azure_exception_mapping.py index 495ca958cf5..249b9349c54 100644 --- a/tests/test_litellm/llms/azure/test_azure_exception_mapping.py +++ b/tests/test_litellm/llms/azure/test_azure_exception_mapping.py @@ -384,4 +384,59 @@ class TestAzureExceptionMapping: model="azure/dall-e-3", original_exception=mock_exception, custom_llm_provider="azure", - ) \ No newline at end of file + ) + + def test_invalid_encrypted_content_error_with_helpful_message(self): + """Test that invalid_encrypted_content errors include helpful guidance + about enabling encrypted_content_affinity.""" + from litellm.exceptions import BadRequestError + + mock_exception = Exception( + "The encrypted content gAAAAABpnW_yEYmSNEyOG... could not be verified. " + "Reason: Encrypted content organization_id did not match the target organization." + ) + mock_exception.body = { + "error": { + "message": "The encrypted content could not be verified.", + "type": "invalid_request_error", + "code": "invalid_encrypted_content", + } + } + mock_response = MagicMock() + mock_response.status_code = 400 + mock_exception.response = mock_response + + with pytest.raises(BadRequestError) as exc_info: + exception_type( + model="azure/gpt-5.1-codex", + original_exception=mock_exception, + custom_llm_provider="azure", + ) + + error = exc_info.value + assert "encrypted_content_affinity" in error.message + assert "enable_pre_call_checks" in error.message + assert "optional_pre_call_checks" in error.message + assert "docs.litellm.ai" in error.message + + def test_openai_invalid_encrypted_content_error(self): + """Test that OpenAI invalid_encrypted_content errors also get helpful guidance.""" + from litellm.exceptions import BadRequestError + + mock_exception = Exception( + "The encrypted content could not be verified." + ) + mock_response = MagicMock() + mock_response.status_code = 400 + mock_exception.response = mock_response + + with pytest.raises(BadRequestError) as exc_info: + exception_type( + model="gpt-5.1-codex", + original_exception=mock_exception, + custom_llm_provider="openai", + ) + + error = exc_info.value + assert "encrypted_content_affinity" in error.message + assert "enable_pre_call_checks" in error.message \ No newline at end of file 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 66177208f4f..6e845e9d050 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 @@ -1,13 +1,18 @@ """ Tests for encrypted_content_affinity pre-call check. -The mechanism works without any cache: -- On response: item IDs for output items with `encrypted_content` are rewritten to - `encitem_{base64("litellm:model_id:{model_id};item_id:{original_id}")}`. -- On routing: `EncryptedContentAffinityCheck` decodes the `encitem_` prefix to extract - `model_id` and pins the request to that deployment. -- Before forwarding: `_restore_encrypted_content_item_ids_in_input` decodes the IDs back - to their original form before sending to the upstream provider. +The mechanism works without any cache and supports two encoding strategies: + +1. **Items with IDs**: item IDs for output items with `encrypted_content` are rewritten to + `encitem_{base64("litellm:model_id:{model_id};item_id:{original_id}")}`. + +2. **Items without IDs** (Codex): encrypted_content itself is wrapped with model_id metadata: + `litellm_enc:{base64("model_id:{model_id}")};{original_encrypted_content}`. + +- On routing: `EncryptedContentAffinityCheck` decodes from either item IDs or wrapped + encrypted_content to extract `model_id` and pins the request to that deployment. +- Before forwarding: `_restore_encrypted_content_item_ids_in_input` decodes IDs and unwraps + encrypted_content back to their original forms before sending to the upstream provider. """ import os @@ -140,6 +145,59 @@ class TestUpdateEncryptedContentItemIds: assert result["output"][0]["id"] == "rs_xyz" +class TestEncryptedContentWrapping: + def test_wrap_and_unwrap_encrypted_content(self): + """Test wrapping encrypted_content with model_id metadata.""" + model_id = "deployment-1" + original_content = "gAAAAABpnW_yEYmSNEyOG_original_encrypted_data" + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + original_content, model_id + ) + assert wrapped.startswith("litellm_enc:") + assert wrapped != original_content + + unwrapped_model_id, unwrapped_content = ( + ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(wrapped) + ) + assert unwrapped_model_id == model_id + assert unwrapped_content == original_content + + def test_unwrap_plain_encrypted_content(self): + """Unwrapping plain encrypted_content returns None for model_id.""" + plain_content = "gAAAAABpnW_yEYmSNEyOG_plain_content" + model_id, content = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( + plain_content + ) + assert model_id is None + assert content == plain_content + + def test_update_response_wraps_encrypted_content_without_id(self): + """Items with encrypted_content but no ID get the content wrapped.""" + model_id = "deployment-1" + response = { + "id": "resp_123", + "output": [ + {"type": "message", "content": []}, + { + "type": "reasoning", + "encrypted_content": "gAAAAABpnW_yEYmSNEyOG_secret", + }, + ], + } + result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( + response, model_id + ) + assert result["output"][0].get("encrypted_content") is None + wrapped = result["output"][1]["encrypted_content"] + assert wrapped.startswith("litellm_enc:") + + model_id_extracted, unwrapped = ( + ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(wrapped) + ) + assert model_id_extracted == model_id + assert unwrapped == "gAAAAABpnW_yEYmSNEyOG_secret" + + class TestRestoreEncryptedContentItemIds: def test_restores_encoded_ids(self): model_id = "deployment-1" @@ -156,6 +214,22 @@ class TestRestoreEncryptedContentItemIds: assert restored[0]["id"] == "msg_abc123" assert restored[1]["id"] == original_id + def test_unwraps_encrypted_content(self): + """Test that wrapped encrypted_content is unwrapped before forwarding.""" + model_id = "deployment-1" + original_content = "gAAAAABpnW_yEYmSNEyOG_original" + wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + original_content, model_id + ) + + request_input = [ + {"type": "reasoning", "encrypted_content": wrapped_content}, + ] + restored = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( + request_input + ) + assert restored[0]["encrypted_content"] == original_content + def test_no_op_for_plain_string_input(self): result = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( "Hello world" @@ -319,8 +393,8 @@ async def test_encrypted_content_affinity_no_effect_on_chat_completions(): @pytest.mark.asyncio async def test_encrypted_content_affinity_bypasses_rpm_limits(): """ - When encrypted content affinity pins to a deployment, RPM limits are bypassed - since the request would fail on any other deployment anyway. + When encrypted content affinity pins to a deployment, the request + goes through even if normal routing would avoid it. """ mock_response_data = { "id": "resp_mock-rpm-test", @@ -347,28 +421,36 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): "litellm_params": { "model": "openai/gpt-5.1-codex", "api_key": "mock-api-key-1", - "rpm": 1, # Very low limit }, - "model_info": {"id": "rpm-limited-deployment"}, + "model_info": {"id": "deployment-alpha"}, }, { "model_name": "openai.gpt-5.1-codex", "litellm_params": { "model": "openai/gpt-5.1-codex", "api_key": "mock-api-key-2", - "rpm": 100, }, - "model_info": {"id": "high-rpm-deployment"}, + "model_info": {"id": "deployment-beta"}, }, ], optional_pre_call_checks=["encrypted_content_affinity"], routing_strategy="usage-based-routing-v2", ) + selected_deployments = [] + + def deterministic_choice(seq): + if len(selected_deployments) == 0: + return seq[0] + return seq[1] if len(seq) > 1 else seq[0] + with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", new_callable=AsyncMock, - ) as mock_post: + ) as mock_post, patch( + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=deterministic_choice, + ): mock_post.return_value = MockResponse(mock_response_data, 200) first_response = await router.aresponses( @@ -376,6 +458,7 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): input="Initial request", ) first_model_id = first_response._hidden_params["model_id"] + selected_deployments.append(first_model_id) # Extract encoded item ID from the first response output encoded_item_id = _extract_encoded_item_id(first_response) @@ -384,7 +467,6 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits(): ) # Follow-up with the encoded item ID — should pin to same deployment - # even if it is at its RPM limit second_response = await router.aresponses( model="openai.gpt-5.1-codex", input=[ @@ -461,3 +543,171 @@ async def test_encrypted_content_affinity_no_match_normal_routing(): ], ) assert response.id is not None + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_with_wrapped_content_no_id(): + """ + Test affinity routing when items have wrapped encrypted_content but no ID. + This simulates Codex client behavior where IDs are omitted. + """ + mock_response_data = { + "id": "resp_mock-wrapped-content", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "openai/gpt-5.1-codex", + "output": [ + { + "type": "reasoning", + "status": "completed", + "encrypted_content": "gAAAAABpnW_yEYmSNEyOG_original_content", + }, + ], + "usage": {"input_tokens": 5, "output_tokens": 10, "total_tokens": 15}, + "error": None, + } + + router = litellm.Router( + model_list=[ + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-1", + }, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": "openai.gpt-5.1-codex", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-2", + }, + "model_info": {"id": "deployment-2"}, + }, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + ) + + selected_deployments = [] + + def deterministic_choice(seq): + if len(selected_deployments) == 0: + return seq[0] + return seq[1] if len(seq) > 1 else seq[0] + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, patch( + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=deterministic_choice, + ): + mock_post.return_value = MockResponse(mock_response_data, 200) + + # First request — goes to deployment-1 + first_response = await router.aresponses( + model="openai.gpt-5.1-codex", + input="Hello, how are you?", + ) + first_model_id = first_response._hidden_params["model_id"] + selected_deployments.append(first_model_id) + + # Extract wrapped encrypted_content from first response + first_item = first_response.output[0] + wrapped_content = ( + first_item.encrypted_content + if hasattr(first_item, "encrypted_content") + else first_item.get("encrypted_content") + ) + assert wrapped_content.startswith("litellm_enc:"), ( + f"Expected wrapped content but got {wrapped_content[:50]}..." + ) + + # Verify we can extract model_id from wrapped content + extracted_model_id, _ = ( + ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( + wrapped_content + ) + ) + assert extracted_model_id == first_model_id + + # Second request: use wrapped encrypted_content WITHOUT an ID (Codex behavior) + second_response = await router.aresponses( + model="openai.gpt-5.1-codex", + input=[ + { + "type": "reasoning", + "encrypted_content": wrapped_content, + }, + ], + ) + second_model_id = second_response._hidden_params["model_id"] + + assert second_model_id == first_model_id, ( + f"Expected affinity to route to {first_model_id}, but got {second_model_id}" + ) + + +def test_encrypted_content_wrapping_preserves_original_content(): + """ + Test that wrapping and unwrapping encrypted_content preserves the original content. + This is critical for streaming responses where content must round-trip correctly. + """ + model_id = "test-deployment-1" + original_encrypted_content = "gAAAAABpnW_yEYmSNEyOG_streaming_test_content_with_special_chars==+/" + + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + original_encrypted_content, model_id + ) + + assert wrapped.startswith("litellm_enc:") + assert wrapped != original_encrypted_content + + extracted_model_id, unwrapped_content = ( + ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(wrapped) + ) + + assert extracted_model_id == model_id + assert unwrapped_content == original_encrypted_content + + +def test_encrypted_content_wrapping_with_multiple_semicolons(): + """ + Test that encrypted_content containing semicolons is handled correctly. + """ + model_id = "deployment-with-semicolons" + original_content = "gAAAAAB;some;content;with;semicolons" + + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + original_content, model_id + ) + + extracted_model_id, unwrapped = ( + ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(wrapped) + ) + + assert extracted_model_id == model_id + assert unwrapped == original_content + + +def test_encrypted_content_wrapping_empty_string(): + """ + Test that empty encrypted_content is handled gracefully. + """ + model_id = "test-deployment" + original_content = "" + + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + original_content, model_id + ) + + assert wrapped.startswith("litellm_enc:") + + extracted_model_id, unwrapped = ( + ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(wrapped) + ) + + assert extracted_model_id == model_id + assert unwrapped == original_content From ca597e18c8a0ca173d45ff1fd56ee39f4d1aa80b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 18:32:41 +0530 Subject: [PATCH 053/147] Fix routing of encrypted content --- litellm/router.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index 652cd68b555..67ccec4c6ef 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1289,6 +1289,11 @@ class Router: ) elif pre_call_check == "enforce_model_rate_limits": _callback = ModelRateLimitingCheck(dual_cache=self.cache) + elif pre_call_check == "encrypted_content_affinity": + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + _callback = EncryptedContentAffinityCheck() if _callback is None: continue From 2f6279d1894ce630afcbf80c1e96231ec06d1e82 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 18:41:15 +0530 Subject: [PATCH 054/147] Fix import issue --- litellm/router.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 67ccec4c6ef..8eb2c417511 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1255,6 +1255,10 @@ 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, + ) + if self.optional_callbacks is None: self.optional_callbacks = [] @@ -1289,11 +1293,6 @@ class Router: ) elif pre_call_check == "enforce_model_rate_limits": _callback = ModelRateLimitingCheck(dual_cache=self.cache) - elif pre_call_check == "encrypted_content_affinity": - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, - ) - _callback = EncryptedContentAffinityCheck() if _callback is None: continue From 5ad0d0367181af342e281e11b8b6c3afc9055e08 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 10:26:03 -0300 Subject: [PATCH 055/147] fix(proxy): encode batch IDs with model info when x-litellm-model header is used When create_batch routes via x-litellm-model header, the response batch_id was returned raw without model routing info. This meant retrieve_batch could not determine which provider/credentials to use, defaulting to "openai" instead of the correct provider (e.g., VLLM). Now encodes batch_id, output_file_id, and error_file_id with model info (same pattern as the model-embedded file_id flow in Scenario 1), so retrieve_batch can decode and route back to the correct provider. --- litellm/proxy/batches_endpoints/endpoints.py | 21 +- .../test_batch_x_litellm_model_encoding.py | 362 ++++++++++++++++++ 2 files changed, 382 insertions(+), 1 deletion(-) create mode 100644 tests/litellm/proxy/test_batch_x_litellm_model_encoding.py diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 1c9ba6cb248..7352e6e2085 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -242,7 +242,26 @@ async def create_batch( # noqa: PLR0915 custom_llm_provider=credentials["custom_llm_provider"], **_create_batch_data # type: ignore ) - + + # Encode response IDs with model info so retrieve_batch + # can route back to the correct provider/credentials. + if response and hasattr(response, "id") and response.id: + response.id = encode_file_id_with_model( + file_id=response.id, + model=model_param, + id_type="batch", + ) + + if hasattr(response, "output_file_id") and response.output_file_id: + response.output_file_id = encode_file_id_with_model( + file_id=response.output_file_id, model=model_param + ) + + if hasattr(response, "error_file_id") and response.error_file_id: + response.error_file_id = encode_file_id_with_model( + file_id=response.error_file_id, model=model_param + ) + verbose_proxy_logger.debug(f"Created batch using model: {model_param}") else: # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py new file mode 100644 index 00000000000..01d6a8aaacd --- /dev/null +++ b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -0,0 +1,362 @@ +""" +Unit tests for batch ID encoding when x-litellm-model header is used. + +Verifies that create_batch encodes response IDs with model info so that +retrieve_batch can route back to the correct provider/credentials. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.proxy.openai_files_endpoints.common_utils import ( + decode_model_from_file_id, + get_original_file_id, +) +from litellm.types.utils import LiteLLMBatch + + +def _make_mock_request(headers: dict) -> MagicMock: + """Create a mock FastAPI Request with the given headers.""" + mock_request = MagicMock() + mock_request.headers = headers + mock_request.query_params = {} + mock_request.url = MagicMock() + mock_request.url.port = 4000 + mock_request.method = "POST" + mock_request.url.path = "/v1/batches" + return mock_request + + +def _make_batch_response( + batch_id: str = "batch_abc123", + input_file_id: str = "file-input456", + output_file_id: str = None, + error_file_id: str = None, + status: str = "validating", +) -> LiteLLMBatch: + """Create a mock LiteLLMBatch response from a provider.""" + return LiteLLMBatch( + id=batch_id, + object="batch", + status=status, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + completion_window="24h", + created_at=1234567890, + output_file_id=output_file_id, + error_file_id=error_file_id, + ) + + +@pytest.mark.asyncio +async def test_create_batch_with_x_litellm_model_encodes_batch_id(): + """ + When x-litellm-model header is provided, create_batch should encode the + response batch_id with model info so retrieve_batch can route correctly. + """ + from litellm.proxy.batches_endpoints.endpoints import create_batch + + model_name = "my-vllm-model" + raw_batch_id = "batch_abc123" + + mock_response = _make_batch_response(batch_id=raw_batch_id) + mock_request = _make_mock_request(headers={"x-litellm-model": model_name}) + mock_fastapi_response = MagicMock() + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.parent_otel_span = None + mock_user_api_key_dict.user_id = "test_user" + + mock_credentials = { + "api_key": "sk-test", + "api_base": "http://vllm:8000", + "custom_llm_provider": "openai", + } + + with ( + patch( + "litellm.proxy.batches_endpoints.endpoints._read_request_body", + new=AsyncMock( + return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"} + ), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processor_cls, + patch( + "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", + return_value=mock_credentials, + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", + ), + patch( + "litellm.acreate_batch", + new=AsyncMock(return_value=mock_response), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.is_known_model", + return_value=False, + ), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), + patch("litellm.proxy.proxy_server.version", "1.0.0"), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + post_call_success_hook=AsyncMock(return_value=mock_response), + update_request_status=AsyncMock(), + ), + ), + ): + # Setup the mock processor to return data and logging obj + mock_processor = MagicMock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=( + {"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + MagicMock(), + ) + ) + mock_processor_cls.return_value = mock_processor + + response = await create_batch( + request=mock_request, + fastapi_response=mock_fastapi_response, + provider=None, + user_api_key_dict=mock_user_api_key_dict, + ) + + # The batch_id should be encoded with model info + assert response.id != raw_batch_id, ( + f"Expected batch_id to be encoded, but got raw ID: {response.id}" + ) + assert response.id.startswith("batch_"), ( + f"Encoded batch_id should keep batch_ prefix, got: {response.id}" + ) + + # Should be decodable back to the original + decoded_model = decode_model_from_file_id(response.id) + assert decoded_model == model_name, ( + f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}" + ) + + original_id = get_original_file_id(response.id) + assert original_id == raw_batch_id, ( + f"Expected original ID '{raw_batch_id}', got: {original_id}" + ) + + +@pytest.mark.asyncio +async def test_create_batch_with_x_litellm_model_encodes_output_and_error_file_ids(): + """ + When a completed batch is returned with output_file_id and error_file_id, + these should also be encoded with model info. + """ + from litellm.proxy.batches_endpoints.endpoints import create_batch + + model_name = "my-vllm-model" + raw_output_file = "file-output789" + raw_error_file = "file-error012" + + mock_response = _make_batch_response( + batch_id="batch_abc123", + output_file_id=raw_output_file, + error_file_id=raw_error_file, + status="completed", + ) + mock_request = _make_mock_request(headers={"x-litellm-model": model_name}) + mock_fastapi_response = MagicMock() + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.parent_otel_span = None + mock_user_api_key_dict.user_id = "test_user" + + mock_credentials = { + "api_key": "sk-test", + "api_base": "http://vllm:8000", + "custom_llm_provider": "openai", + } + + with ( + patch( + "litellm.proxy.batches_endpoints.endpoints._read_request_body", + new=AsyncMock( + return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"} + ), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processor_cls, + patch( + "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", + return_value=mock_credentials, + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", + ), + patch( + "litellm.acreate_batch", + new=AsyncMock(return_value=mock_response), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.is_known_model", + return_value=False, + ), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), + patch("litellm.proxy.proxy_server.version", "1.0.0"), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + post_call_success_hook=AsyncMock(return_value=mock_response), + update_request_status=AsyncMock(), + ), + ), + ): + mock_processor = MagicMock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=( + {"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + MagicMock(), + ) + ) + mock_processor_cls.return_value = mock_processor + + response = await create_batch( + request=mock_request, + fastapi_response=mock_fastapi_response, + provider=None, + user_api_key_dict=mock_user_api_key_dict, + ) + + # output_file_id should be encoded + assert decode_model_from_file_id(response.output_file_id) == model_name + assert get_original_file_id(response.output_file_id) == raw_output_file + + # error_file_id should be encoded + assert decode_model_from_file_id(response.error_file_id) == model_name + assert get_original_file_id(response.error_file_id) == raw_error_file + + +@pytest.mark.asyncio +async def test_create_batch_without_x_litellm_model_returns_raw_ids(): + """ + Without x-litellm-model header, create_batch should NOT encode batch IDs + (falls through to Scenario 3 / custom_llm_provider fallback). + """ + from litellm.proxy.batches_endpoints.endpoints import create_batch + + raw_batch_id = "batch_abc123" + mock_response = _make_batch_response(batch_id=raw_batch_id) + mock_request = _make_mock_request(headers={}) + mock_fastapi_response = MagicMock() + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.parent_otel_span = None + mock_user_api_key_dict.user_id = "test_user" + + with ( + patch( + "litellm.proxy.batches_endpoints.endpoints._read_request_body", + new=AsyncMock( + return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"} + ), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processor_cls, + patch( + "litellm.acreate_batch", + new=AsyncMock(return_value=mock_response), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.is_known_model", + return_value=False, + ), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), + patch("litellm.proxy.proxy_server.version", "1.0.0"), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + post_call_success_hook=AsyncMock(return_value=mock_response), + update_request_status=AsyncMock(), + ), + ), + ): + mock_processor = MagicMock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=( + {"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + MagicMock(), + ) + ) + mock_processor_cls.return_value = mock_processor + + response = await create_batch( + request=mock_request, + fastapi_response=mock_fastapi_response, + provider=None, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Without x-litellm-model, the batch_id should remain raw + assert response.id == raw_batch_id + assert decode_model_from_file_id(response.id) is None + + +class TestBatchIdRoundTripWithRetrieve: + """ + Tests that batch IDs encoded during create_batch can be decoded + correctly during retrieve_batch (Scenario 1: model_from_id). + """ + + def test_encoded_batch_id_is_decoded_for_retrieve(self): + """ + Simulates the full round-trip: create encodes the ID, + retrieve decodes it to get the model and original batch_id. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + ) + + model_name = "my-vllm-model" + raw_batch_id = "batch_vllm_12345" + + # What create_batch does: + encoded_id = encode_file_id_with_model( + file_id=raw_batch_id, model=model_name, id_type="batch" + ) + + # What retrieve_batch does: + decoded_model = decode_model_from_file_id(encoded_id) + original_id = get_original_file_id(encoded_id) + + assert decoded_model == model_name + assert original_id == raw_batch_id + + def test_vllm_style_batch_id_roundtrip(self): + """ + VLLM may return batch IDs in various formats. + Verify round-trip works for common patterns. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + ) + + test_cases = [ + ("batch_abc123", "vllm-llama3"), + ("batch_67890", "openai/llama-3-8b"), + ("batch_some-uuid-here", "my-custom-vllm"), + ] + + for raw_id, model in test_cases: + encoded = encode_file_id_with_model( + file_id=raw_id, model=model, id_type="batch" + ) + assert encoded.startswith("batch_") + assert decode_model_from_file_id(encoded) == model + assert get_original_file_id(encoded) == raw_id From 3426b905cedd66d59506174efb04d8c53566b53b Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Mar 2026 10:44:47 -0300 Subject: [PATCH 056/147] Update tests/litellm/proxy/test_batch_x_litellm_model_encoding.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/litellm/proxy/test_batch_x_litellm_model_encoding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py index 01d6a8aaacd..062443d9425 100644 --- a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -33,8 +33,8 @@ def _make_mock_request(headers: dict) -> MagicMock: def _make_batch_response( batch_id: str = "batch_abc123", input_file_id: str = "file-input456", - output_file_id: str = None, - error_file_id: str = None, + output_file_id: Optional[str] = None, + error_file_id: Optional[str] = None, status: str = "validating", ) -> LiteLLMBatch: """Create a mock LiteLLMBatch response from a provider.""" From 7d664f0c096a7f85befc17b22d80c8cbf9097a93 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Mar 2026 10:45:00 -0300 Subject: [PATCH 057/147] Update tests/litellm/proxy/test_batch_x_litellm_model_encoding.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/litellm/proxy/test_batch_x_litellm_model_encoding.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py index 062443d9425..5150b57568b 100644 --- a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -5,7 +5,6 @@ Verifies that create_batch encodes response IDs with model info so that retrieve_batch can route back to the correct provider/credentials. """ -import json from unittest.mock import AsyncMock, MagicMock, patch import pytest From 096edface5b7772eccf398113a57b26941c00457 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Mar 2026 10:46:14 -0300 Subject: [PATCH 058/147] Update litellm/proxy/batches_endpoints/endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/batches_endpoints/endpoints.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 7352e6e2085..58c8e2d4d0c 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -257,11 +257,16 @@ async def create_batch( # noqa: PLR0915 file_id=response.output_file_id, model=model_param ) - if hasattr(response, "error_file_id") and response.error_file_id: + if hasattr(response, "error_file_id") and response.error_file_id: response.error_file_id = encode_file_id_with_model( file_id=response.error_file_id, model=model_param ) + if hasattr(response, "input_file_id") and response.input_file_id: + response.input_file_id = encode_file_id_with_model( + file_id=response.input_file_id, model=model_param + ) + verbose_proxy_logger.debug(f"Created batch using model: {model_param}") else: # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) From d66f8bc15d8f22c52d81785ab047ff65208da35f Mon Sep 17 00:00:00 2001 From: Varad Khonde Date: Tue, 3 Mar 2026 19:17:28 +0530 Subject: [PATCH 059/147] feat(togetherai): add support for togetherai/Qwen3.5-397B-A17B model --- litellm/model_prices_and_context_window_backup.json | 12 ++++++++++++ model_prices_and_context_window.json | 12 ++++++++++++ 2 files changed, 24 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d4c5b476af6..e476d0c6b2a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29397,6 +29397,18 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/Qwen/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4934f11d456..cf55ddf6e25 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29632,6 +29632,18 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/Qwen/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", From 9463de0c6629354dbc616b1a1beb394f27a45b80 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 10:51:48 -0300 Subject: [PATCH 060/147] fix: correct indentation from commit suggestions and add missing Optional import --- litellm/proxy/batches_endpoints/endpoints.py | 2 +- tests/litellm/proxy/test_batch_x_litellm_model_encoding.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 58c8e2d4d0c..cdee69f2b30 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -257,7 +257,7 @@ async def create_batch( # noqa: PLR0915 file_id=response.output_file_id, model=model_param ) - if hasattr(response, "error_file_id") and response.error_file_id: + if hasattr(response, "error_file_id") and response.error_file_id: response.error_file_id = encode_file_id_with_model( file_id=response.error_file_id, model=model_param ) diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py index 5150b57568b..521a3632dcb 100644 --- a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -5,6 +5,7 @@ Verifies that create_batch encodes response IDs with model info so that retrieve_batch can route back to the correct provider/credentials. """ +from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest From 6c0387d170a2caab45393222bcc32e769a2c5124 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 19:31:55 +0530 Subject: [PATCH 061/147] Add support for Attaching knowledge base to model via UI --- .../src/components/add_model/AddModelForm.tsx | 1 + .../add_model/advanced_settings.test.tsx | 3 + .../add_model/advanced_settings.tsx | 30 +++++++++ .../src/components/model_info_view.tsx | 61 +++++++++++++++++++ 4 files changed, 95 insertions(+) diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 59ac63cffe6..2b3f23a35ae 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -358,6 +358,7 @@ const AddModelForm: React.FC = ({ teams={teams} guardrailsList={guardrailsList || []} tagsList={tagsList || {}} + accessToken={accessToken || ""} /> )} diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx index 6515c67c292..9fe36e13998 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx @@ -13,6 +13,7 @@ describe("AdvancedSettings", () => { setShowAdvancedSettings={() => {}} guardrailsList={[]} tagsList={{}} + accessToken="test-token" />, ); }); @@ -24,6 +25,7 @@ describe("AdvancedSettings", () => { setShowAdvancedSettings={() => {}} guardrailsList={[]} tagsList={{}} + accessToken="test-token" />, ); fireEvent.click(getByText("Advanced Settings")); @@ -39,6 +41,7 @@ describe("AdvancedSettings", () => { setShowAdvancedSettings={() => {}} guardrailsList={[]} tagsList={{}} + accessToken="test-token" />, ); act(() => { diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx index c9f5ef8a4b1..8ae90c1cbbc 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx @@ -6,6 +6,7 @@ import TextArea from "antd/es/input/TextArea"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Team } from "../key_team_helpers/key_list"; import CacheControlSettings from "./cache_control_settings"; +import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; import { Tag } from "../tag_management/types"; import { formItemValidateJSON } from "../../utils/textUtils"; const { Link } = Typography; @@ -16,6 +17,7 @@ interface AdvancedSettingsProps { teams?: Team[] | null; guardrailsList: string[]; tagsList: Record; + accessToken: string; } const AdvancedSettings: React.FC = ({ @@ -24,6 +26,7 @@ const AdvancedSettings: React.FC = ({ teams, guardrailsList, tagsList, + accessToken, }) => { const [form] = Form.useForm(); const [customPricing, setCustomPricing] = React.useState(false); @@ -109,6 +112,33 @@ const AdvancedSettings: React.FC = ({ + + Attached Knowledge Bases (RAG){" "} + + e.stopPropagation()} + > + + + + + } + name="vector_store_ids" + className="mt-4" + help="Select vector stores to attach. Requests to this model will automatically use these for RAG. Set up vector stores in Tools > Vector Stores." + > + {}} + accessToken={accessToken} + placeholder="Select knowledge bases (optional)" + /> + + diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index e2fc8caa21c..40c1a3a386a 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -17,6 +17,7 @@ import { Button as TremorButton, } from "@tremor/react"; import { Button, Form, Input, Modal, Select, Tooltip } from "antd"; +import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import { CheckIcon, CopyIcon } from "lucide-react"; import { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../utils/dataUtils"; @@ -245,6 +246,11 @@ export default function ModelInfoView({ if (values.guardrails) { updatedLitellmParams.guardrails = values.guardrails; } + if (values.vector_store_ids !== undefined) { + updatedLitellmParams.vector_store_ids = Array.isArray(values.vector_store_ids) + ? values.vector_store_ids + : []; + } // Handle cache control settings if (values.cache_control && values.cache_control_injection_points?.length > 0) { @@ -606,6 +612,9 @@ export default function ModelInfoView({ guardrails: Array.isArray(localModelData.litellm_params?.guardrails) ? localModelData.litellm_params.guardrails : [], + vector_store_ids: Array.isArray(localModelData.litellm_params?.vector_store_ids) + ? localModelData.litellm_params.vector_store_ids + : [], tags: Array.isArray(localModelData.litellm_params?.tags) ? localModelData.litellm_params.tags : [], health_check_model: isWildcardModel ? localModelData.model_info?.health_check_model : null, litellm_extra_params: JSON.stringify(localModelData.litellm_params || {}, null, 2), @@ -883,6 +892,58 @@ export default function ModelInfoView({ )} +
+ + Attached Knowledge Bases (RAG) + + e.stopPropagation()} + > + + + + + {isEditing ? ( + + {}} + accessToken={accessToken || ""} + placeholder="Select knowledge bases (optional)" + /> + + ) : ( +
+ {localModelData.litellm_params?.vector_store_ids ? ( + Array.isArray(localModelData.litellm_params.vector_store_ids) ? ( + localModelData.litellm_params.vector_store_ids.length > 0 ? ( +
+ {localModelData.litellm_params.vector_store_ids.map( + (vsId: string, index: number) => ( + + {vsId} + + ) + )} +
+ ) : ( + "No knowledge bases attached" + ) + ) : ( + String(localModelData.litellm_params.vector_store_ids) + ) + ) : ( + "Not Set" + )} +
+ )} +
+
Tags {isEditing ? ( From 24ec7f882f5e3036d47544f8ef96f899fb4606aa Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 19:35:26 +0530 Subject: [PATCH 062/147] Revert "feat(togetherai): add support for togetherai/Qwen3.5-397B-A17B model" --- litellm/model_prices_and_context_window_backup.json | 12 ------------ model_prices_and_context_window.json | 12 ------------ 2 files changed, 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e476d0c6b2a..d4c5b476af6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29397,18 +29397,6 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "together_ai/Qwen/Qwen/Qwen3.5-397B-A17B": { - "input_cost_per_token": 6e-07, - "litellm_provider": "together_ai", - "max_input_tokens": 262144, - "mode": "chat", - "output_cost_per_token": 3.6e-06, - "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cf55ddf6e25..4934f11d456 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29632,18 +29632,6 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "together_ai/Qwen/Qwen/Qwen3.5-397B-A17B": { - "input_cost_per_token": 6e-07, - "litellm_provider": "together_ai", - "max_input_tokens": 262144, - "mode": "chat", - "output_cost_per_token": 3.6e-06, - "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", From 9fd4c00b064e6358fdb5202daa417b3cfb482796 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 11:06:38 -0300 Subject: [PATCH 063/147] fix(proxy): re-encode response IDs in retrieve_batch for model-based routing The provider returns raw IDs in the retrieve response (output_file_id, error_file_id). These need to be encoded with model info so the client can use them for subsequent file download calls through the proxy. --- litellm/proxy/batches_endpoints/endpoints.py | 26 ++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index cdee69f2b30..d3e983c5ad5 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -464,8 +464,30 @@ async def retrieve_batch( # noqa: PLR0915 custom_llm_provider=credentials["custom_llm_provider"], **data # type: ignore ) - - + + # Re-encode response IDs so the client always sees encoded IDs. + # The provider returns raw IDs (e.g. output_file_id, error_file_id) + # which the client needs encoded to route future file downloads. + if response and hasattr(response, "id") and response.id: + response.id = encode_file_id_with_model( + file_id=response.id, model=model_from_id, id_type="batch", + ) + + if hasattr(response, "output_file_id") and response.output_file_id: + response.output_file_id = encode_file_id_with_model( + file_id=response.output_file_id, model=model_from_id + ) + + if hasattr(response, "error_file_id") and response.error_file_id: + response.error_file_id = encode_file_id_with_model( + file_id=response.error_file_id, model=model_from_id + ) + + if hasattr(response, "input_file_id") and response.input_file_id: + response.input_file_id = encode_file_id_with_model( + file_id=response.input_file_id, model=model_from_id + ) + verbose_proxy_logger.debug( f"Retrieved batch using model: {model_from_id}, original_id: {original_batch_id}" ) From 7506fd0426a70c0387fd811d44fdd06c1fc0006b Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 11:22:43 -0300 Subject: [PATCH 064/147] fix(proxy): re-encode response IDs in cancel_batch for model-based routing --- litellm/proxy/batches_endpoints/endpoints.py | 23 +++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index d3e983c5ad5..7ab50e321ac 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -855,7 +855,28 @@ async def cancel_batch( custom_llm_provider=credentials["custom_llm_provider"], **data # type: ignore ) - + + # Re-encode response IDs so the client always sees encoded IDs. + if response and hasattr(response, "id") and response.id: + response.id = encode_file_id_with_model( + file_id=response.id, model=model_from_id, id_type="batch", + ) + + if hasattr(response, "output_file_id") and response.output_file_id: + response.output_file_id = encode_file_id_with_model( + file_id=response.output_file_id, model=model_from_id + ) + + if hasattr(response, "error_file_id") and response.error_file_id: + response.error_file_id = encode_file_id_with_model( + file_id=response.error_file_id, model=model_from_id + ) + + if hasattr(response, "input_file_id") and response.input_file_id: + response.input_file_id = encode_file_id_with_model( + file_id=response.input_file_id, model=model_from_id + ) + verbose_proxy_logger.debug( f"Cancelled batch using model: {model_from_id}, original_id: {original_batch_id}" ) From 59bde4a81a92c9e7b8f48ca2e32311a735a51a7d Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 11:38:50 -0300 Subject: [PATCH 065/147] refactor(proxy): extract encode_batch_response_ids helper and fix list_batches encoding Extract duplicated batch ID encoding logic into a shared helper encode_batch_response_ids() in common_utils.py. Use it in create_batch, retrieve_batch, and cancel_batch. Also add encoding to list_batches when x-litellm-model is used. --- litellm/proxy/batches_endpoints/endpoints.py | 77 +++---------------- .../openai_files_endpoints/common_utils.py | 16 ++++ 2 files changed, 27 insertions(+), 66 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 7ab50e321ac..ae99d59c631 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, decode_model_from_file_id, + encode_batch_response_ids, encode_file_id_with_model, get_batch_from_database, get_credentials_for_model, @@ -243,29 +244,7 @@ async def create_batch( # noqa: PLR0915 **_create_batch_data # type: ignore ) - # Encode response IDs with model info so retrieve_batch - # can route back to the correct provider/credentials. - if response and hasattr(response, "id") and response.id: - response.id = encode_file_id_with_model( - file_id=response.id, - model=model_param, - id_type="batch", - ) - - if hasattr(response, "output_file_id") and response.output_file_id: - response.output_file_id = encode_file_id_with_model( - file_id=response.output_file_id, model=model_param - ) - - if hasattr(response, "error_file_id") and response.error_file_id: - response.error_file_id = encode_file_id_with_model( - file_id=response.error_file_id, model=model_param - ) - - if hasattr(response, "input_file_id") and response.input_file_id: - response.input_file_id = encode_file_id_with_model( - file_id=response.input_file_id, model=model_param - ) + encode_batch_response_ids(response, model=model_param) verbose_proxy_logger.debug(f"Created batch using model: {model_param}") else: @@ -465,28 +444,7 @@ async def retrieve_batch( # noqa: PLR0915 **data # type: ignore ) - # Re-encode response IDs so the client always sees encoded IDs. - # The provider returns raw IDs (e.g. output_file_id, error_file_id) - # which the client needs encoded to route future file downloads. - if response and hasattr(response, "id") and response.id: - response.id = encode_file_id_with_model( - file_id=response.id, model=model_from_id, id_type="batch", - ) - - if hasattr(response, "output_file_id") and response.output_file_id: - response.output_file_id = encode_file_id_with_model( - file_id=response.output_file_id, model=model_from_id - ) - - if hasattr(response, "error_file_id") and response.error_file_id: - response.error_file_id = encode_file_id_with_model( - file_id=response.error_file_id, model=model_from_id - ) - - if hasattr(response, "input_file_id") and response.input_file_id: - response.input_file_id = encode_file_id_with_model( - file_id=response.input_file_id, model=model_from_id - ) + encode_batch_response_ids(response, model=model_from_id) verbose_proxy_logger.debug( f"Retrieved batch using model: {model_from_id}, original_id: {original_batch_id}" @@ -679,7 +637,13 @@ async def list_batches( limit=limit, **data # type: ignore ) - + + # Encode batch IDs in the list response so clients can use + # them for retrieve/cancel/file downloads through the proxy. + if response and hasattr(response, "data") and response.data: + for batch in response.data: + encode_batch_response_ids(batch, model=model_param) + verbose_proxy_logger.debug(f"Listed batches using model: {model_param}") # SCENARIO 2 (alternative): target_model_names based routing @@ -856,26 +820,7 @@ async def cancel_batch( **data # type: ignore ) - # Re-encode response IDs so the client always sees encoded IDs. - if response and hasattr(response, "id") and response.id: - response.id = encode_file_id_with_model( - file_id=response.id, model=model_from_id, id_type="batch", - ) - - if hasattr(response, "output_file_id") and response.output_file_id: - response.output_file_id = encode_file_id_with_model( - file_id=response.output_file_id, model=model_from_id - ) - - if hasattr(response, "error_file_id") and response.error_file_id: - response.error_file_id = encode_file_id_with_model( - file_id=response.error_file_id, model=model_from_id - ) - - if hasattr(response, "input_file_id") and response.input_file_id: - response.input_file_id = encode_file_id_with_model( - file_id=response.input_file_id, model=model_from_id - ) + encode_batch_response_ids(response, model=model_from_id) verbose_proxy_logger.debug( f"Cancelled batch using model: {model_from_id}, original_id: {original_batch_id}" diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index ceaf3c7550e..343ea119672 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -129,6 +129,22 @@ def encode_file_id_with_model( return f"{prefix}{encoded_b64}" +def encode_batch_response_ids(response, model: str) -> None: + """Encode all IDs in a batch response with model routing info (in-place).""" + if not response or not hasattr(response, "id") or not response.id: + return + response.id = encode_file_id_with_model( + file_id=response.id, model=model, id_type="batch" + ) + for attr in ("output_file_id", "error_file_id", "input_file_id"): + if hasattr(response, attr) and getattr(response, attr): + setattr( + response, + attr, + encode_file_id_with_model(file_id=getattr(response, attr), model=model), + ) + + def decode_model_from_file_id(encoded_id: str) -> Optional[str]: """ Extract model name from an encoded file/batch ID. From 6d535e56395782de8841497aa228729e16354872 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 3 Mar 2026 20:14:18 +0530 Subject: [PATCH 066/147] fix(proxy): allow custom auth routes to bypass route authorization checks Custom user-added routes (e.g. /ldap/ngs/ready) used with Depends(user_api_key_auth) were being rejected as admin-only after _run_post_custom_auth_checks was introduced in commit 14badde13c. The route authorization check in common_checks is designed for LiteLLM's own management routes. Custom auth flows that add their own routes should be trusted since the custom auth function already validated the request. Budget and expiry checks still run. Add skip_route_check parameter to common_checks() and pass skip_route_check=True from _run_post_custom_auth_checks() to skip route authorization while preserving budget/team/model checks. Regression test added: test_common_checks_skip_route_check_for_custom_auth Co-Authored-By: Claude Haiku 4.5 --- litellm/proxy/auth/auth_checks.py | 28 +++++----- litellm/proxy/auth/user_api_key_auth.py | 1 + .../proxy/auth/test_auth_checks.py | 52 +++++++++++++++++++ 3 files changed, 69 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index ef6b0ac462c..91ac58215ab 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -234,6 +234,7 @@ async def common_checks( request: Request, skip_budget_checks: bool = False, project_object: Optional[LiteLLM_ProjectTableCachedObj] = None, + skip_route_check: bool = False, ) -> bool: """ Common checks across jwt + key-based auth. @@ -453,18 +454,21 @@ async def common_checks( user_object=user_object, route=route, request_body=request_body ) - token_team = getattr(valid_token, "team_id", None) - token_type: Literal["ui", "api"] = ( - "ui" if token_team is not None and token_team == "litellm-dashboard" else "api" - ) - _is_route_allowed = _is_allowed_route( - route=route, - token_type=token_type, - user_obj=user_object, - request=request, - request_data=request_body, - valid_token=valid_token, - ) + if not skip_route_check: + token_team = getattr(valid_token, "team_id", None) + token_type: Literal["ui", "api"] = ( + "ui" + if token_team is not None and token_team == "litellm-dashboard" + else "api" + ) + _is_route_allowed = _is_allowed_route( + route=route, + token_type=token_type, + user_obj=user_object, + request=request, + request_data=request_body, + valid_token=valid_token, + ) # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store await vector_store_access_check( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a2705ceb7da..1d575fb5131 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1766,6 +1766,7 @@ async def _run_post_custom_auth_checks( valid_token=valid_token, skip_budget_checks=False, project_object=_project_obj, + skip_route_check=True, ) return valid_token diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 4cdca2d0617..501c2285d1e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1519,3 +1519,55 @@ async def test_get_fuzzy_user_object_case_insensitive_email(): assert call_args.kwargs["where"]["user_email"]["equals"] == "test@example.com" assert call_args.kwargs["where"]["user_email"]["mode"] == "insensitive" assert call_args.kwargs["include"] == {"organization_memberships": True} + + +@pytest.mark.asyncio +async def test_common_checks_skip_route_check_for_custom_auth(): + """ + Test that custom routes (e.g. /ldap/ngs/ready) pass common_checks when + skip_route_check=True, which is the case for custom auth flows. + + Regression test for: custom user-added routes being rejected as admin-only + after _run_post_custom_auth_checks was introduced. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + mock_request = MagicMock(spec=Request) + valid_token = UserAPIKeyAuth(token="test-token") + + # Without skip_route_check, a custom route with unknown user should fail + with pytest.raises(Exception): + await common_checks( + request_body={}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/ldap/ngs/ready", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=mock_request, + skip_route_check=False, + ) + + # With skip_route_check=True (custom auth path), the same route should pass + result = await common_checks( + request_body={}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/ldap/ngs/ready", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=mock_request, + skip_route_check=True, + ) + + assert result is True From 75518c3ca73626b457a7697fea456a0cabc1a341 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 12:03:40 -0300 Subject: [PATCH 067/147] feat(models): add zai/glm-5 and zai/glm-5-code to model cost map Add native ZhipuAI GLM-5 and GLM-5-Code model entries with pricing from docs.z.ai/guides/overview/pricing. --- model_prices_and_context_window.json | 30 ++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4c0694db371..ac0e9dfaa7a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33811,6 +33811,36 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "zai/glm-5": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, + "zai/glm-5-code": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 5e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.7": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 1.1e-07, From b83373d29c43667d108e087b95a5cadc9262dba9 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Tue, 3 Mar 2026 11:01:49 -0500 Subject: [PATCH 068/147] Managed batches - Address PR bot comments from #22464 --- litellm/files/main.py | 2 +- litellm/llms/vertex_ai/batches/handler.py | 5 +- .../llms/vertex_ai/files/transformation.py | 5 +- .../test_file_retrieve_provider_routing.py | 127 ++++++++++++++++++ .../test_vertex_ai_files_transformation.py | 48 ++++++- 5 files changed, 181 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py diff --git a/litellm/files/main.py b/litellm/files/main.py index 66d3a97468d..f0a8112fbdf 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -336,7 +336,7 @@ async def afile_retrieve( @client def file_retrieve( file_id: str, - custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "manus"] = "openai", + custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index ba3b5fb7a2c..5f1fefca963 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -115,9 +115,10 @@ class VertexAIBatchPrediction(VertexLLM): data=json.dumps(vertex_batch_request), ) except httpx.HTTPStatusError as e: - error_body = e.response.text if hasattr(e, 'response') else "N/A" + error_body = e.response.text litellm.verbose_logger.error( - f"Vertex AI batch create failed: status={e.response.status_code}, body={error_body[:1000]}" + "Vertex AI batch create failed: status=%s, body=%s", + e.response.status_code, error_body[:1000], ) raise if response.status_code != 200: diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index f0493cd6be9..bf3ed5e6ac9 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -408,10 +408,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): file_id = "deleted" if hasattr(raw_response, "request") and raw_response.request: url = str(raw_response.request.url) - if "/o/" in url: + if "/b/" in url and "/o/" in url: import urllib.parse + bucket_part = url.split("/b/")[-1].split("/o/")[0] encoded_name = url.split("/o/")[-1].split("?")[0] - file_id = f"gs://{urllib.parse.unquote(encoded_name)}" + file_id = f"gs://{bucket_part}/{urllib.parse.unquote(encoded_name)}" return FileDeleted(id=file_id, deleted=True, object="file") def transform_list_files_request( diff --git a/tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py b/tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py new file mode 100644 index 00000000000..68d5e2035f7 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py @@ -0,0 +1,127 @@ +""" +Tests for Fix 1: file_retrieve Literal type was missing 'vertex_ai' and 'gemini', +causing a type mismatch when afile_retrieve delegated to the sync function. +""" + +import pytest +from unittest.mock import MagicMock, patch + +from litellm.files.main import file_retrieve + + +class TestFileRetrieveProviderRouting: + """ + Verify that file_retrieve accepts 'vertex_ai' and 'gemini' providers and + routes them through ProviderConfigManager / base_llm_http_handler. + """ + + def _make_mock_file_object(self): + mock = MagicMock() + mock.model_dump.return_value = { + "id": "gs://my-bucket/file.jsonl", + "object": "file", + "bytes": 1024, + "created_at": 0, + "filename": "file.jsonl", + "purpose": "batch", + "status": "processed", + } + return mock + + def test_should_route_vertex_ai_through_provider_config(self): + """ + Regression: file_retrieve Literal type was missing 'vertex_ai', + so passing custom_llm_provider='vertex_ai' would fail type-checking + and potentially cause a routing failure at runtime. + """ + mock_file = self._make_mock_file_object() + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ) as mock_retrieve: + result = file_retrieve( + file_id="gs://my-bucket/file.jsonl", + custom_llm_provider="vertex_ai", + ) + + mock_retrieve.assert_called_once() + assert result is not None + + def test_should_route_gemini_through_provider_config(self): + """ + Regression: file_retrieve Literal type was also missing 'gemini'. + """ + mock_file = self._make_mock_file_object() + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ) as mock_retrieve: + result = file_retrieve( + file_id="some-gemini-file-id", + custom_llm_provider="gemini", + ) + + mock_retrieve.assert_called_once() + assert result is not None + + def test_should_pass_file_id_to_handler_for_vertex_ai(self): + """Verify the file_id is forwarded correctly to the underlying handler.""" + mock_file = self._make_mock_file_object() + expected_file_id = "gs://my-bucket/path/to/file.jsonl" + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ) as mock_retrieve: + file_retrieve( + file_id=expected_file_id, + custom_llm_provider="vertex_ai", + ) + + call_kwargs = mock_retrieve.call_args.kwargs + assert call_kwargs.get("file_id") == expected_file_id + + def test_should_not_raise_bad_request_for_vertex_ai(self): + """ + Before the fix, vertex_ai fell through to the else-branch which raised + BadRequestError. Verify it no longer does. + """ + import litellm + + mock_file = self._make_mock_file_object() + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ): + try: + file_retrieve( + file_id="gs://my-bucket/file.jsonl", + custom_llm_provider="vertex_ai", + ) + except litellm.exceptions.BadRequestError as e: + pytest.fail( + f"file_retrieve raised BadRequestError for vertex_ai: {e}" + ) + + def test_should_not_raise_bad_request_for_gemini(self): + """Same as above but for 'gemini'.""" + import litellm + + mock_file = self._make_mock_file_object() + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ): + try: + file_retrieve( + file_id="some-file-id", + custom_llm_provider="gemini", + ) + except litellm.exceptions.BadRequestError as e: + pytest.fail( + f"file_retrieve raised BadRequestError for gemini: {e}" + ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 6f1d753484d..598ad255aca 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -167,7 +167,7 @@ class TestTransformDeleteFile: assert isinstance(result, FileDeleted) assert result.deleted is True assert result.object == "file" - assert "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" in result.id + assert result.id == "gs://my-bucket/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" def test_should_fallback_to_deleted_id_when_no_request(self, config): raw_response = MagicMock(spec=httpx.Response) @@ -182,3 +182,49 @@ class TestTransformDeleteFile: assert isinstance(result, FileDeleted) assert result.id == "deleted" assert result.deleted is True + + def test_should_include_bucket_name_in_reconstructed_delete_id(self, config): + """ + Regression: the old code split on /o/ only, dropping the bucket from + the reconstructed gs:// URI. e.g. gs://path/to/file instead of + gs://my-bucket/path/to/file. + """ + raw_response = MagicMock(spec=httpx.Response) + mock_request = MagicMock() + encoded_object = urllib.parse.quote("path/to/file.jsonl", safe="") + mock_request.url = ( + f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_object}" + ) + raw_response.request = mock_request + + result = config.transform_delete_file_response( + raw_response=raw_response, + logging_obj=MagicMock(), + litellm_params={}, + ) + + assert result.id == "gs://my-bucket/path/to/file.jsonl" + + def test_should_include_bucket_in_nested_object_path(self, config): + """Verify bucket extraction works with deeply nested GCS object paths.""" + raw_response = MagicMock(spec=httpx.Response) + mock_request = MagicMock() + encoded_object = urllib.parse.quote( + "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123", + safe="", + ) + mock_request.url = ( + f"https://storage.googleapis.com/storage/v1/b/prod-bucket/o/{encoded_object}" + ) + raw_response.request = mock_request + + result = config.transform_delete_file_response( + raw_response=raw_response, + logging_obj=MagicMock(), + litellm_params={}, + ) + + assert result.id == ( + "gs://prod-bucket/litellm-vertex-files/publishers/google/" + "models/gemini-2.0-flash-001/abc-123" + ) From c3fe4634b62599cea71856d351571a681f3a8e12 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 22:06:16 +0530 Subject: [PATCH 069/147] Add correct pricing for gemini 3.1 flash lite --- ...odel_prices_and_context_window_backup.json | 33 ++++++++++++------- model_prices_and_context_window.json | 33 ++++++++++++------- 2 files changed, 42 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 42b4f0f7762..8f5fd26e3e8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14335,17 +14335,18 @@ "supports_web_search": true }, "gemini-3.1-flash-lite-preview": { - "cache_read_input_token_cost": 2.5e-09, - "input_cost_per_audio_token": 2.5e-07, + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", @@ -14368,6 +14369,8 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -17138,17 +17141,18 @@ "supports_service_tier": true }, "gemini/gemini-3.1-flash-lite-preview": { - "cache_read_input_token_cost": 2.5e-09, - "input_cost_per_audio_token": 2.5e-07, + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", @@ -17172,6 +17176,8 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32400,17 +32406,18 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, "vertex_ai/gemini-3.1-flash-lite-preview": { - "cache_read_input_token_cost": 2.5e-09, - "input_cost_per_audio_token": 2.5e-07, + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", @@ -32433,6 +32440,8 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 42b4f0f7762..8f5fd26e3e8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14335,17 +14335,18 @@ "supports_web_search": true }, "gemini-3.1-flash-lite-preview": { - "cache_read_input_token_cost": 2.5e-09, - "input_cost_per_audio_token": 2.5e-07, + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", @@ -14368,6 +14369,8 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -17138,17 +17141,18 @@ "supports_service_tier": true }, "gemini/gemini-3.1-flash-lite-preview": { - "cache_read_input_token_cost": 2.5e-09, - "input_cost_per_audio_token": 2.5e-07, + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", @@ -17172,6 +17176,8 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32400,17 +32406,18 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, "vertex_ai/gemini-3.1-flash-lite-preview": { - "cache_read_input_token_cost": 2.5e-09, - "input_cost_per_audio_token": 2.5e-07, + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "max_images_per_prompt": 3000, "max_input_tokens": 1048576, - "max_output_tokens": 65535, + "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 65535, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", @@ -32433,6 +32440,8 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, From 9d06106af077b0996c154c43c014b5db0b509094 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 3 Mar 2026 22:22:57 +0530 Subject: [PATCH 070/147] Fix gemini-3.1-flash-lite-preview for streaming --- litellm/model_prices_and_context_window_backup.json | 3 ++- model_prices_and_context_window.json | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8f5fd26e3e8..5de764c5cec 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14382,7 +14382,8 @@ "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_native_streaming": true }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8f5fd26e3e8..5de764c5cec 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14382,7 +14382,8 @@ "supports_url_context": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "supports_native_streaming": true }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, From 22e682b1e89c1c1795b45910b2cea988f1525124 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 3 Mar 2026 22:39:06 +0530 Subject: [PATCH 071/147] feat: guardrail-mode-default-list --- .../docs/proxy/guardrails/quick_start.md | 40 +++++++++++++++++++ litellm/integrations/custom_guardrail.py | 18 ++++++++- litellm/types/guardrails.py | 2 +- 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/proxy/guardrails/quick_start.md b/docs/my-website/docs/proxy/guardrails/quick_start.md index e5a90f74a8a..eb56c27f876 100644 --- a/docs/my-website/docs/proxy/guardrails/quick_start.md +++ b/docs/my-website/docs/proxy/guardrails/quick_start.md @@ -499,6 +499,11 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ Run guardrails based on the user-agent header. This is useful for running pre-call checks on OpenWebUI but only masking in logs for Claude CLI. +`default` can be a single mode string or a list of modes. + + + + ```yaml model_list: - model_name: gpt-3.5-turbo @@ -519,6 +524,32 @@ guardrails: default_on: true # run on every request ``` + + + +```yaml +model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: gpt-3.5-turbo + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "guardrails_ai-guard" + litellm_params: + guardrail: guardrails_ai + guard_name: "pii_detect" + mode: + tags: + "User-Agent: claude-cli": "logging_only" + default: ["pre_call", "post_call"] # Run on both pre and post call when no tags match + api_base: os.environ/GUARDRAILS_AI_API_BASE + default_on: true +``` + + + + ### ✨ Model-level Guardrails @@ -640,13 +671,22 @@ guardrails: Mode Specification +`default` accepts either a single string or a list of strings. + ```python from litellm.types.guardrails import Mode +# Single default mode mode = Mode( tags={"User-Agent: claude-cli": "logging_only"}, default="logging_only" ) + +# Multiple default modes +mode = Mode( + tags={"User-Agent: claude-cli": "logging_only"}, + default=["pre_call", "post_call"] +) ``` ### `guardrails` Request Parameter diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 5d11fd68475..3fd179bb411 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -235,8 +235,13 @@ class CustomGuardrail(CustomLogger): list(event_hook.tags.values()), supported_event_hooks ) if event_hook.default: + default_list = ( + event_hook.default + if isinstance(event_hook.default, list) + else [event_hook.default] + ) _validate_event_hook_list_is_in_supported_event_hooks( - [event_hook.default], supported_event_hooks + default_list, supported_event_hooks ) elif isinstance(event_hook, GuardrailEventHooks): if event_hook not in supported_event_hooks: @@ -461,7 +466,16 @@ class CustomGuardrail(CustomLogger): if isinstance(self.event_hook, list): return event_type.value in self.event_hook if isinstance(self.event_hook, Mode): - return event_type.value in self.event_hook.tags.values() + if event_type.value in self.event_hook.tags.values(): + return True + if self.event_hook.default: + default_list = ( + self.event_hook.default + if isinstance(self.event_hook.default, list) + else [self.event_hook.default] + ) + return event_type.value in default_list + return False return self.event_hook == event_type.value def get_guardrail_dynamic_request_body_params(self, request_data: dict) -> dict: diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index a68c4e2f762..dc95ed3314a 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -718,7 +718,7 @@ class BaseLitellmParams( class Mode(BaseModel): tags: Dict[str, str] = Field(description="Tags for the guardrail mode") - default: Optional[str] = Field( + default: Optional[Union[str, List[str]]] = Field( default=None, description="Default mode when no tags match" ) From b44755db96fa54c6d2e4c38960ea20cbac0a93fc Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 3 Mar 2026 22:50:34 +0530 Subject: [PATCH 072/147] fix(proxy): make common_checks opt-in for custom auth via custom_auth_run_common_checks Replaces the skip_route_check approach from PR #22662 with a configurable opt-in flag. By default, common_checks() is not run for custom auth flows, preserving backwards compatibility with pre-#22164 behavior. Users who want budget/team/route enforcement on custom auth can enable it: general_settings: custom_auth_run_common_checks: true Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/auth/auth_checks.py | 28 +++---- litellm/proxy/auth/user_api_key_auth.py | 32 ++++---- .../proxy/auth/test_auth_checks.py | 77 +++++++++---------- .../auth/test_custom_auth_end_user_budget.py | 21 ++++- 4 files changed, 86 insertions(+), 72 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 91ac58215ab..ef6b0ac462c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -234,7 +234,6 @@ async def common_checks( request: Request, skip_budget_checks: bool = False, project_object: Optional[LiteLLM_ProjectTableCachedObj] = None, - skip_route_check: bool = False, ) -> bool: """ Common checks across jwt + key-based auth. @@ -454,21 +453,18 @@ async def common_checks( user_object=user_object, route=route, request_body=request_body ) - if not skip_route_check: - token_team = getattr(valid_token, "team_id", None) - token_type: Literal["ui", "api"] = ( - "ui" - if token_team is not None and token_team == "litellm-dashboard" - else "api" - ) - _is_route_allowed = _is_allowed_route( - route=route, - token_type=token_type, - user_obj=user_object, - request=request, - request_data=request_body, - valid_token=valid_token, - ) + token_team = getattr(valid_token, "team_id", None) + token_type: Literal["ui", "api"] = ( + "ui" if token_team is not None and token_team == "litellm-dashboard" else "api" + ) + _is_route_allowed = _is_allowed_route( + route=route, + token_type=token_type, + user_obj=user_object, + request=request, + request_data=request_body, + valid_token=valid_token, + ) # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store await vector_store_access_check( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 1d575fb5131..7b52f6bb96d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1752,21 +1752,21 @@ async def _run_post_custom_auth_checks( if _project_obj is not None: valid_token.project_metadata = _project_obj.metadata - _ = await common_checks( - request=request, - request_body=request_data, - team_object=_team_obj, - user_object=user_object, - end_user_object=end_user_object, - general_settings=general_settings, - global_proxy_spend=None, - route=route, - llm_router=llm_router, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - skip_budget_checks=False, - project_object=_project_obj, - skip_route_check=True, - ) + if general_settings.get("custom_auth_run_common_checks", False): + _ = await common_checks( + request=request, + request_body=request_data, + team_object=_team_obj, + user_object=user_object, + end_user_object=end_user_object, + general_settings=general_settings, + global_proxy_spend=None, + route=route, + llm_router=llm_router, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + skip_budget_checks=False, + project_object=_project_obj, + ) return valid_token diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 501c2285d1e..ff1bc5b2581 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1522,52 +1522,51 @@ async def test_get_fuzzy_user_object_case_insensitive_email(): @pytest.mark.asyncio -async def test_common_checks_skip_route_check_for_custom_auth(): +async def test_custom_auth_common_checks_opt_in(): """ - Test that custom routes (e.g. /ldap/ngs/ready) pass common_checks when - skip_route_check=True, which is the case for custom auth flows. + Test that _run_post_custom_auth_checks only runs common_checks when + custom_auth_run_common_checks is explicitly set to True in general_settings. - Regression test for: custom user-added routes being rejected as admin-only - after _run_post_custom_auth_checks was introduced. + By default (False), common_checks is skipped for backwards compatibility + with custom auth flows that existed before PR #22164. """ - from fastapi import Request + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks - from litellm.proxy.auth.auth_checks import common_checks - - mock_request = MagicMock(spec=Request) valid_token = UserAPIKeyAuth(token="test-token") + mock_request = MagicMock() - # Without skip_route_check, a custom route with unknown user should fail - with pytest.raises(Exception): - await common_checks( - request_body={}, - team_object=None, - user_object=None, - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/ldap/ngs/ready", - llm_router=None, - proxy_logging_obj=MagicMock(), + # Default (no flag) — common_checks should NOT be called + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ): + mock_common.return_value = True + result = await _run_post_custom_auth_checks( valid_token=valid_token, request=mock_request, - skip_route_check=False, + request_data={}, + route="/ldap/ngs/ready", + parent_otel_span=None, ) + mock_common.assert_not_called() - # With skip_route_check=True (custom auth path), the same route should pass - result = await common_checks( - request_body={}, - team_object=None, - user_object=None, - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/ldap/ngs/ready", - llm_router=None, - proxy_logging_obj=MagicMock(), - valid_token=valid_token, - request=mock_request, - skip_route_check=True, - ) - - assert result is True + # With flag=True — common_checks SHOULD be called + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ): + mock_common.return_value = True + result = await _run_post_custom_auth_checks( + valid_token=valid_token, + request=mock_request, + request_data={}, + route="/chat/completions", + parent_otel_span=None, + ) + mock_common.assert_called_once() diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index 73a97188424..18816dcec4a 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -10,9 +10,10 @@ from litellm.proxy._types import UserAPIKeyAuth @pytest.mark.asyncio async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): - # Test backwards compatibility + # Test backwards compatibility — common_checks only runs when opt-in flag is set valid_token = UserAPIKeyAuth(token="test_token") + # Default: common_checks should NOT be called with patch( "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock ) as mock_common: @@ -26,6 +27,24 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): ) assert result.token == "test_token" assert getattr(result, "end_user_id", None) is None + mock_common.assert_not_awaited() + + # With opt-in flag: common_checks SHOULD be called + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ): + mock_common.return_value = True + result = await _run_post_custom_auth_checks( + valid_token=valid_token, + request=None, + request_data={}, + route="/v1/chat/completions", + parent_otel_span=None, + ) + assert result.token == "test_token" mock_common.assert_awaited_once() From 1fe2e92d3272bcf1c7d21f401efd8b8a32bbd475 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 14:34:32 -0300 Subject: [PATCH 073/147] fix(main): forward enable_json_schema_validation to acompletion_with_mcp The parameter was declared in completion() signature but not passed to acompletion_with_mcp, causing per-request JSON schema validation to silently fall back to the global default when MCP tools are present. --- litellm/main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/main.py b/litellm/main.py index 362e5b8b263..75353c1b070 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1170,6 +1170,7 @@ def completion( # type: ignore # noqa: PLR0915 thinking=thinking, web_search_options=web_search_options, shared_session=shared_session, + enable_json_schema_validation=enable_json_schema_validation, **kwargs, ) api_base = kwargs.get("api_base", None) From 76e3dba0f88929a82c673b7c586cc4aba8b8f34a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 09:41:45 -0800 Subject: [PATCH 074/147] fix mcp server created_at and updated_at timestamps being overwritten with current time - Add created_at field to MCPServer type (was missing) - Map created_at from LiteLLM_MCPServerTable in build_mcp_server_from_table() - Use server.created_at and server.updated_at instead of datetime.now() in _build_mcp_server_table() and health check table builder - Add regression tests to verify timestamps are preserved through round-trip conversions Co-Authored-By: Claude Sonnet 4.6 --- .../mcp_server/mcp_server_manager.py | 11 ++- .../types/mcp_server/mcp_server_manager.py | 1 + .../mcp_server/test_mcp_server_manager.py | 86 +++++++++++++++++++ 3 files changed, 92 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index da29c7804a1..b7c013e9f20 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -642,6 +642,7 @@ class MCPServerManager: available_on_public_internet=bool( getattr(mcp_server, "available_on_public_internet", True) ), + created_at=getattr(mcp_server, "created_at", None), updated_at=getattr(mcp_server, "updated_at", None), ) return new_server @@ -2540,8 +2541,8 @@ class MCPServerManager: url=server.url, transport=server.transport, auth_type=server.auth_type, - created_at=datetime.now(), - updated_at=datetime.now(), + created_at=server.created_at, + updated_at=server.updated_at, teams=[], mcp_access_groups=server.access_groups or [], allowed_tools=server.allowed_tools or [], @@ -2620,8 +2621,6 @@ class MCPServerManager: return list_mcp_servers def _build_mcp_server_table(self, server: MCPServer) -> LiteLLM_MCPServerTable: - from datetime import datetime - return LiteLLM_MCPServerTable( server_id=server.server_id, server_name=server.server_name, @@ -2633,8 +2632,8 @@ class MCPServerManager: spec_path=server.spec_path, transport=server.transport, auth_type=server.auth_type, - created_at=datetime.now(), - updated_at=datetime.now(), + created_at=server.created_at, + updated_at=server.updated_at, teams=[], mcp_access_groups=server.access_groups or [], allowed_tools=server.allowed_tools or [], diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 69b34a25a21..cabac6b9d51 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -53,6 +53,7 @@ class MCPServer(BaseModel): access_groups: Optional[List[str]] = None allow_all_keys: bool = False available_on_public_internet: bool = True + created_at: Optional[datetime] = None updated_at: Optional[datetime] = None model_config = ConfigDict(arbitrary_types_allowed=True) 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 c105052479d..acc76221cbb 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 @@ -2307,5 +2307,91 @@ class TestMCPServerManager: assert resolved_server.server_name == "test_server" # server_name matches +class TestMCPServerTimestamps: + """Regression tests: created_at/updated_at must be preserved, not overwritten with datetime.now().""" + + @pytest.mark.asyncio + async def test_build_mcp_server_from_table_preserves_timestamps(self): + """build_mcp_server_from_table must carry created_at and updated_at into MCPServer.""" + manager = MCPServerManager() + + created = datetime(2024, 1, 15, 10, 0, 0) + updated = datetime(2024, 6, 20, 12, 30, 0) + + table_record = LiteLLM_MCPServerTable( + server_id="ts-server-1", + server_name="ts_server", + url="https://example.com/mcp", + transport=MCPTransport.http, + created_at=created, + updated_at=updated, + ) + + mcp_server = await manager.build_mcp_server_from_table(table_record) + + assert mcp_server.created_at == created + assert mcp_server.updated_at == updated + + def test_build_mcp_server_table_preserves_timestamps(self): + """_build_mcp_server_table must use the MCPServer's stored timestamps, not datetime.now().""" + manager = MCPServerManager() + + created = datetime(2024, 1, 15, 10, 0, 0) + updated = datetime(2024, 6, 20, 12, 30, 0) + + server = MCPServer( + server_id="ts-server-2", + name="ts_server", + url="https://example.com/mcp", + transport=MCPTransport.http, + created_at=created, + updated_at=updated, + ) + + table = manager._build_mcp_server_table(server) + + assert table.created_at == created + assert table.updated_at == updated + + def test_build_mcp_server_table_none_timestamps_when_not_set(self): + """_build_mcp_server_table must return None timestamps when not set on MCPServer.""" + manager = MCPServerManager() + + server = MCPServer( + server_id="ts-server-3", + name="ts_server", + url="https://example.com/mcp", + transport=MCPTransport.http, + ) + + table = manager._build_mcp_server_table(server) + + assert table.created_at is None + assert table.updated_at is None + + @pytest.mark.asyncio + async def test_round_trip_timestamps_preserved(self): + """Timestamps survive the full round-trip: LiteLLM_MCPServerTable -> MCPServer -> LiteLLM_MCPServerTable.""" + manager = MCPServerManager() + + created = datetime(2023, 3, 10, 8, 0, 0) + updated = datetime(2023, 9, 5, 16, 45, 0) + + table_record = LiteLLM_MCPServerTable( + server_id="ts-server-4", + server_name="ts_server_rt", + url="https://example.com/mcp", + transport=MCPTransport.http, + created_at=created, + updated_at=updated, + ) + + mcp_server = await manager.build_mcp_server_from_table(table_record) + rebuilt_table = manager._build_mcp_server_table(mcp_server) + + assert rebuilt_table.created_at == created + assert rebuilt_table.updated_at == updated + + if __name__ == "__main__": pytest.main([__file__]) From f6264a9c0fde055be2cdb2b84a2dee6c32c8885a Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 14:50:41 -0300 Subject: [PATCH 075/147] docs(openrouter): add image edit documentation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add OpenRouter image edit docs to both the provider page and the main image_edits reference page, including supported models, parameter mappings (size→aspect_ratio, quality→image_size), usage examples, proxy configuration, and a note about 4K quality model support. --- docs/my-website/docs/image_edits.md | 71 +++++++++++++++- docs/my-website/docs/providers/openrouter.md | 87 ++++++++++++++++++++ 2 files changed, 157 insertions(+), 1 deletion(-) diff --git a/docs/my-website/docs/image_edits.md b/docs/my-website/docs/image_edits.md index a8438334542..f1cfc0ed8e9 100644 --- a/docs/my-website/docs/image_edits.md +++ b/docs/my-website/docs/image_edits.md @@ -16,7 +16,7 @@ LiteLLM provides image editing functionality that maps to OpenAI's `/images/edit | Supported operations | Create image edits | Single and multiple images supported | | Supported LiteLLM SDK Versions | 1.63.8+ | Gemini support requires 1.79.3+ | | Supported LiteLLM Proxy Versions | 1.71.1+ | Gemini support requires 1.79.3+ | -| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **Stability AI**, **AWS Bedrock (Stability)** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. Stability AI and Bedrock Stability support various image editing operations. | +| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **OpenRouter**, **Stability AI**, **AWS Bedrock (Stability)** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. OpenRouter routes image edits through chat completions. Stability AI and Bedrock Stability support various image editing operations. | #### ⚡️See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/) @@ -244,6 +244,47 @@ response = litellm.image_edit( print(response) ``` + + + + +#### Basic Image Edit +```python showLineNumbers title="OpenRouter Image Edit" +import os +from litellm import image_edit + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=open("original_image.png", "rb"), + prompt="Add aurora borealis to the night sky", +) + +print(response) +``` + +#### Multiple Images Edit +```python showLineNumbers title="OpenRouter Multiple Images Edit" +import os +from litellm import image_edit + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=[ + open("scene.png", "rb"), + open("style_reference.png", "rb"), + ], + prompt="Blend the reference style into the scene", + size="1536x1024", # mapped to aspect_ratio 3:2 + quality="high", # mapped to image_size 4K +) + +print(response) +``` + @@ -398,6 +439,34 @@ curl -X POST "http://0.0.0.0:4000/v1/images/edits" \ -F "size=1024x1024" ``` + + + + +1. Add the OpenRouter image edit model to your `config.yaml`: +```yaml showLineNumbers title="OpenRouter Proxy Configuration" +model_list: + - model_name: openrouter-image-edit + litellm_params: + model: openrouter/google/gemini-2.5-flash-image + api_key: os.environ/OPENROUTER_API_KEY +``` + +2. Start the LiteLLM proxy server: +```bash showLineNumbers title="Start LiteLLM Proxy Server" +litellm --config /path/to/config.yaml +``` + +3. Make an image edit request: +```bash showLineNumbers title="OpenRouter Proxy Image Edit" +curl -X POST "http://0.0.0.0:4000/v1/images/edits" \ + -H "Authorization: Bearer " \ + -F "model=openrouter-image-edit" \ + -F "image=@original_image.png" \ + -F "prompt=Make the sky a vibrant purple sunset" \ + -F "size=1024x1024" +``` + diff --git a/docs/my-website/docs/providers/openrouter.md b/docs/my-website/docs/providers/openrouter.md index 38eb998c98b..4c79c41cfd5 100644 --- a/docs/my-website/docs/providers/openrouter.md +++ b/docs/my-website/docs/providers/openrouter.md @@ -210,3 +210,90 @@ response = image_generation( # Cost is available in the response metadata print(f"Request cost: ${response._hidden_params['additional_headers']['llm_provider-x-litellm-response-cost']}") ``` + +## Image Edit + +OpenRouter supports image editing through select models like Google Gemini image models. LiteLLM routes image edit requests to OpenRouter's chat completions endpoint with the source image sent as a base64 data URL and `modalities: ["image", "text"]`. + +### Supported Models + +| Model | Description | +|-------|-------------| +| `openrouter/google/gemini-2.5-flash-image` | Gemini 2.5 Flash with image editing | + +See all available image models on [OpenRouter's model list](https://openrouter.ai/models?modality=image). + +### Supported Parameters + +| Parameter | OpenRouter Mapping | Notes | +|-----------|--------------------|-------| +| `size` | `image_config.aspect_ratio` | `1024x1024` → `1:1`, `1536x1024` → `3:2`, `1024x1536` → `2:3`, `1792x1024` → `16:9`, `1024x1792` → `9:16` | +| `quality` | `image_config.image_size` | `low`/`standard` → `1K`, `medium` → `2K`, `high`/`hd` → `4K` | +| `n` | `n` | Number of images | + +:::note +`quality=high` (4K) is only supported by `google/gemini-3-pro-image-preview` and `google/gemini-3.1-flash-image-preview`. The `google/gemini-2.5-flash-image` model supports up to `medium` (2K). +::: + +### Usage + +```python +from litellm import image_edit +import os + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +# Basic image edit +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=open("original_image.png", "rb"), + prompt="Make the sky a vibrant purple sunset", +) + +print(response) +``` + +### Advanced Usage with Parameters + +```python +from litellm import image_edit +import os + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +# Edit with size and quality parameters +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=open("photo.png", "rb"), + prompt="Add northern lights to the sky", + size="1536x1024", # Maps to aspect_ratio 3:2 + quality="high", # Maps to image_size 4K +) + +# Access the edited image +image_data = response.data[0] +if image_data.b64_json: + import base64 + with open("edited.png", "wb") as f: + f.write(base64.b64decode(image_data.b64_json)) +``` + +### Multiple Images Edit + +```python +from litellm import image_edit +import os + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=[ + open("scene.png", "rb"), + open("style_reference.png", "rb"), + ], + prompt="Blend the reference style into the scene", +) + +print(response) +``` From 909e3ce6c9c0bcc3bbf0ef320766fd3e5d712726 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 14:54:11 -0300 Subject: [PATCH 076/147] test: create fresh ModelResponse per test to avoid shared mutable state --- .../test_json_schema_validation.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/tests/litellm/litellm_core_utils/test_json_schema_validation.py b/tests/litellm/litellm_core_utils/test_json_schema_validation.py index 859a1238a61..f798db6fb43 100644 --- a/tests/litellm/litellm_core_utils/test_json_schema_validation.py +++ b/tests/litellm/litellm_core_utils/test_json_schema_validation.py @@ -46,11 +46,8 @@ STRICT_SCHEMA = { }, } -# Response that does NOT match the schema (wrong field names) -INVALID_RESPONSE = _make_response({"name": "test", "age": 25}) - -# Response that matches the schema -VALID_RESPONSE = _make_response({"title": "Inception", "rating": 9}) +INVALID_CONTENT = {"name": "test", "age": 25} # Does NOT match the schema +VALID_CONTENT = {"title": "Inception", "rating": 9} # Matches the schema @pytest.fixture(autouse=True) @@ -70,7 +67,7 @@ class TestPerRequestJsonSchemaValidation: litellm.enable_json_schema_validation = False # Should NOT raise even though response doesn't match schema post_call_processing( - INVALID_RESPONSE, + _make_response(INVALID_CONTENT), "test-model", {"response_format": STRICT_SCHEMA}, _mock_completion, @@ -82,7 +79,7 @@ class TestPerRequestJsonSchemaValidation: litellm.enable_json_schema_validation = False with pytest.raises(litellm.JSONSchemaValidationError): post_call_processing( - INVALID_RESPONSE, + _make_response(INVALID_CONTENT), "test-model", { "response_format": STRICT_SCHEMA, @@ -97,7 +94,7 @@ class TestPerRequestJsonSchemaValidation: litellm.enable_json_schema_validation = True # Should NOT raise because per-request says False post_call_processing( - INVALID_RESPONSE, + _make_response(INVALID_CONTENT), "test-model", { "response_format": STRICT_SCHEMA, @@ -112,7 +109,7 @@ class TestPerRequestJsonSchemaValidation: litellm.enable_json_schema_validation = True with pytest.raises(litellm.JSONSchemaValidationError): post_call_processing( - INVALID_RESPONSE, + _make_response(INVALID_CONTENT), "test-model", {"response_format": STRICT_SCHEMA}, _mock_completion, @@ -122,7 +119,7 @@ class TestPerRequestJsonSchemaValidation: def test_valid_response_passes_with_per_request_on(self): """Per-request ON + valid response -> no error raised.""" post_call_processing( - VALID_RESPONSE, + _make_response(VALID_CONTENT), "test-model", { "response_format": STRICT_SCHEMA, From fba19f089a1950c957f6b68b8ddf7d7a252af479 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 10 Feb 2026 23:33:54 -0300 Subject: [PATCH 077/147] refactor: reduce code duplication in files/main.py with credential helpers Extract repeated OpenAI and Azure credential resolution logic into _get_openai_credentials() and _get_azure_credentials() helper functions, reducing ~270 lines of duplicated code across 5 file operations. Also removes dead Vertex AI code path in create_file that was unreachable since ProviderConfigManager.get_provider_files_config() handles it first. --- litellm/files/main.py | 374 ++++++++++++------------------------------ 1 file changed, 103 insertions(+), 271 deletions(-) diff --git a/litellm/files/main.py b/litellm/files/main.py index 78e41bb5a68..879e6fe69c0 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -11,7 +11,7 @@ import os import time import uuid as uuid_module from functools import partial -from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast +from typing import Any, Coroutine, Dict, Literal, NamedTuple, Optional, Union, cast import httpx @@ -57,6 +57,68 @@ anthropic_files_instance = AnthropicFilesHandler() ################################################# +class OpenAICredentials(NamedTuple): + api_base: Optional[str] + api_key: Optional[str] + organization: Optional[str] + + +class AzureCredentials(NamedTuple): + api_base: Optional[str] + api_key: Optional[str] + api_version: Optional[str] + + +def _get_openai_credentials( + optional_params: GenericLiteLLMParams, +) -> OpenAICredentials: + """Resolve OpenAI credentials from optional_params, litellm globals, and env vars.""" + api_base = ( + optional_params.api_base + or litellm.api_base + or os.getenv("OPENAI_BASE_URL") + or os.getenv("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + organization = ( + optional_params.organization + or litellm.organization + or os.getenv("OPENAI_ORGANIZATION", None) + or None + ) + api_key = ( + optional_params.api_key + or litellm.api_key + or litellm.openai_key + or os.getenv("OPENAI_API_KEY") + ) + return OpenAICredentials(api_base=api_base, api_key=api_key, organization=organization) + + +def _get_azure_credentials( + optional_params: GenericLiteLLMParams, +) -> AzureCredentials: + """Resolve Azure credentials from optional_params, litellm globals, and env vars.""" + api_base = ( + optional_params.api_base + or litellm.api_base + or get_secret_str("AZURE_API_BASE") + ) + api_version = ( + optional_params.api_version + or litellm.api_version + or get_secret_str("AZURE_API_VERSION") + ) + api_key = ( + optional_params.api_key + or litellm.api_key + or litellm.azure_key + or get_secret_str("AZURE_OPENAI_API_KEY") + or get_secret_str("AZURE_API_KEY") + ) + return AzureCredentials(api_base=api_base, api_key=api_key, api_version=api_version) + + @client async def acreate_file( file: FileTypes, @@ -185,95 +247,28 @@ def create_file( timeout=timeout, ) elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there - api_base = ( - optional_params.api_base - or litellm.api_base - or os.getenv("OPENAI_BASE_URL") - or os.getenv("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 - ) - # set API KEY - api_key = ( - optional_params.api_key - or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") - ) - + openai_creds = _get_openai_credentials(optional_params) response = openai_files_instance.create_file( _is_async=_is_async, - api_base=api_base, - api_key=api_key, + api_base=openai_creds.api_base, + api_key=openai_creds.api_key, timeout=timeout, max_retries=optional_params.max_retries, - organization=organization, + organization=openai_creds.organization, create_file_data=_create_file_request, ) elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore - - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) # type: ignore - - extra_body = optional_params.get("extra_body", {}) - if extra_body is not None: - extra_body.pop("azure_ad_token", None) - else: - get_secret_str("AZURE_AD_TOKEN") # type: ignore - + azure_creds = _get_azure_credentials(optional_params) response = azure_files_instance.create_file( _is_async=_is_async, - api_base=api_base, - api_key=api_key, - api_version=api_version, + api_base=azure_creds.api_base, + api_key=azure_creds.api_key, + api_version=azure_creds.api_version, timeout=timeout, max_retries=optional_params.max_retries, create_file_data=_create_file_request, litellm_params=litellm_params_dict, ) - elif custom_llm_provider == "vertex_ai": - api_base = optional_params.api_base or "" - vertex_ai_project = ( - optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") - ) - vertex_ai_location = ( - optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str( - "VERTEXAI_CREDENTIALS" - ) - - response = vertex_ai_files_instance.create_file( - _is_async=_is_async, - api_base=api_base, - vertex_project=vertex_ai_project, - vertex_location=vertex_ai_location, - vertex_credentials=vertex_credentials, - timeout=timeout, - max_retries=optional_params.max_retries, - create_file_data=_create_file_request, - ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus'] are supported.".format( @@ -367,64 +362,23 @@ def file_retrieve( _is_async = kwargs.pop("is_async", False) is True if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there - api_base = ( - optional_params.api_base - or litellm.api_base - or os.getenv("OPENAI_BASE_URL") - or os.getenv("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 - ) - # set API KEY - api_key = ( - optional_params.api_key - or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") - ) - + openai_creds = _get_openai_credentials(optional_params) response = openai_files_instance.retrieve_file( file_id=file_id, _is_async=_is_async, - api_base=api_base, - api_key=api_key, + api_base=openai_creds.api_base, + api_key=openai_creds.api_key, timeout=timeout, max_retries=optional_params.max_retries, - organization=organization, + organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore - - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) # type: ignore - - extra_body = optional_params.get("extra_body", {}) - if extra_body is not None: - extra_body.pop("azure_ad_token", None) - else: - get_secret_str("AZURE_AD_TOKEN") # type: ignore - + azure_creds = _get_azure_credentials(optional_params) response = azure_files_instance.retrieve_file( _is_async=_is_async, - api_base=api_base, - api_key=api_key, - api_version=api_version, + api_base=azure_creds.api_base, + api_key=azure_creds.api_key, + api_version=azure_creds.api_version, timeout=timeout, max_retries=optional_params.max_retries, file_id=file_id, @@ -576,63 +530,23 @@ def file_delete( timeout = 600.0 _is_async = kwargs.pop("is_async", False) is True if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there - api_base = ( - optional_params.api_base - or litellm.api_base - or os.getenv("OPENAI_BASE_URL") - or os.getenv("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 - ) - # set API KEY - api_key = ( - optional_params.api_key - or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") - ) + openai_creds = _get_openai_credentials(optional_params) response = openai_files_instance.delete_file( file_id=file_id, _is_async=_is_async, - api_base=api_base, - api_key=api_key, + api_base=openai_creds.api_base, + api_key=openai_creds.api_key, timeout=timeout, max_retries=optional_params.max_retries, - organization=organization, + organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore - - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) # type: ignore - - extra_body = optional_params.get("extra_body", {}) - if extra_body is not None: - extra_body.pop("azure_ad_token", None) - else: - get_secret_str("AZURE_AD_TOKEN") # type: ignore - + azure_creds = _get_azure_credentials(optional_params) response = azure_files_instance.delete_file( _is_async=_is_async, - api_base=api_base, - api_key=api_key, - api_version=api_version, + api_base=azure_creds.api_base, + api_key=azure_creds.api_key, + api_version=azure_creds.api_version, timeout=timeout, max_retries=optional_params.max_retries, file_id=file_id, @@ -815,64 +729,23 @@ def file_list( ) return response elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there - api_base = ( - optional_params.api_base - or litellm.api_base - or os.getenv("OPENAI_BASE_URL") - or os.getenv("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 - ) - # set API KEY - api_key = ( - optional_params.api_key - or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") - ) - + openai_creds = _get_openai_credentials(optional_params) response = openai_files_instance.list_files( purpose=purpose, _is_async=_is_async, - api_base=api_base, - api_key=api_key, + api_base=openai_creds.api_base, + api_key=openai_creds.api_key, timeout=timeout, max_retries=optional_params.max_retries, - organization=organization, + organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore - - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) # type: ignore - - extra_body = optional_params.get("extra_body", {}) - if extra_body is not None: - extra_body.pop("azure_ad_token", None) - else: - get_secret_str("AZURE_AD_TOKEN") # type: ignore - + azure_creds = _get_azure_credentials(optional_params) response = azure_files_instance.list_files( _is_async=_is_async, - api_base=api_base, - api_key=api_key, - api_version=api_version, + api_base=azure_creds.api_base, + api_key=azure_creds.api_key, + api_version=azure_creds.api_version, timeout=timeout, max_retries=optional_params.max_retries, purpose=purpose, @@ -1003,64 +876,23 @@ def file_content( return response if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there - api_base = ( - optional_params.api_base - or litellm.api_base - or os.getenv("OPENAI_BASE_URL") - or os.getenv("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 - ) - # set API KEY - api_key = ( - optional_params.api_key - or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") - ) - + openai_creds = _get_openai_credentials(optional_params) response = openai_files_instance.file_content( _is_async=_is_async, file_content_request=_file_content_request, - api_base=api_base, - api_key=api_key, + api_base=openai_creds.api_base, + api_key=openai_creds.api_key, timeout=timeout, max_retries=optional_params.max_retries, - organization=organization, + organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore - - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) # type: ignore - - extra_body = optional_params.get("extra_body", {}) - if extra_body is not None: - extra_body.pop("azure_ad_token", None) - else: - get_secret_str("AZURE_AD_TOKEN") # type: ignore - + azure_creds = _get_azure_credentials(optional_params) response = azure_files_instance.file_content( _is_async=_is_async, - api_base=api_base, - api_key=api_key, - api_version=api_version, + api_base=azure_creds.api_base, + api_key=azure_creds.api_key, + api_version=azure_creds.api_version, timeout=timeout, max_retries=optional_params.max_retries, file_content_request=_file_content_request, From ead74dff114454ca5eba36c33bd09e36fedc3e4c Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 10 Feb 2026 23:41:53 -0300 Subject: [PATCH 078/147] refactor: move credential helpers to provider common_utils modules Move get_openai_credentials() to litellm/llms/openai/common_utils.py and get_azure_credentials() to litellm/llms/azure/common_utils.py so they can be reused by batches/main.py and other modules. Signatures now take individual params instead of GenericLiteLLMParams. --- litellm/files/main.py | 153 ++++++++++++++-------------- litellm/llms/azure/common_utils.py | 38 ++++++- litellm/llms/openai/common_utils.py | 39 ++++++- 3 files changed, 154 insertions(+), 76 deletions(-) diff --git a/litellm/files/main.py b/litellm/files/main.py index 879e6fe69c0..f626f9b3466 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -7,11 +7,10 @@ https://platform.openai.com/docs/api-reference/files import asyncio import contextvars -import os import time import uuid as uuid_module from functools import partial -from typing import Any, Coroutine, Dict, Literal, NamedTuple, Optional, Union, cast +from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast import httpx @@ -20,10 +19,12 @@ from litellm import get_secret_str from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.files.handler import AnthropicFilesHandler +from litellm.llms.azure.common_utils import get_azure_credentials from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI from litellm.llms.bedrock.files.handler import BedrockFilesHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai.openai import FileDeleted, FileObject, OpenAIFilesAPI from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler from litellm.types.llms.openai import ( @@ -57,68 +58,6 @@ anthropic_files_instance = AnthropicFilesHandler() ################################################# -class OpenAICredentials(NamedTuple): - api_base: Optional[str] - api_key: Optional[str] - organization: Optional[str] - - -class AzureCredentials(NamedTuple): - api_base: Optional[str] - api_key: Optional[str] - api_version: Optional[str] - - -def _get_openai_credentials( - optional_params: GenericLiteLLMParams, -) -> OpenAICredentials: - """Resolve OpenAI credentials from optional_params, litellm globals, and env vars.""" - api_base = ( - optional_params.api_base - or litellm.api_base - or os.getenv("OPENAI_BASE_URL") - or os.getenv("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None - ) - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") - ) - return OpenAICredentials(api_base=api_base, api_key=api_key, organization=organization) - - -def _get_azure_credentials( - optional_params: GenericLiteLLMParams, -) -> AzureCredentials: - """Resolve Azure credentials from optional_params, litellm globals, and env vars.""" - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) - return AzureCredentials(api_base=api_base, api_key=api_key, api_version=api_version) - - @client async def acreate_file( file: FileTypes, @@ -247,7 +186,11 @@ def create_file( timeout=timeout, ) elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = _get_openai_credentials(optional_params) + openai_creds = get_openai_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + organization=optional_params.organization, + ) response = openai_files_instance.create_file( _is_async=_is_async, api_base=openai_creds.api_base, @@ -258,7 +201,11 @@ def create_file( create_file_data=_create_file_request, ) elif custom_llm_provider == "azure": - azure_creds = _get_azure_credentials(optional_params) + azure_creds = get_azure_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + api_version=optional_params.api_version, + ) response = azure_files_instance.create_file( _is_async=_is_async, api_base=azure_creds.api_base, @@ -269,6 +216,32 @@ def create_file( create_file_data=_create_file_request, litellm_params=litellm_params_dict, ) + elif custom_llm_provider == "vertex_ai": + api_base = optional_params.api_base or "" + vertex_ai_project = ( + optional_params.vertex_project + or litellm.vertex_project + or get_secret_str("VERTEXAI_PROJECT") + ) + vertex_ai_location = ( + optional_params.vertex_location + or litellm.vertex_location + or get_secret_str("VERTEXAI_LOCATION") + ) + vertex_credentials = optional_params.vertex_credentials or get_secret_str( + "VERTEXAI_CREDENTIALS" + ) + + response = vertex_ai_files_instance.create_file( + _is_async=_is_async, + api_base=api_base, + vertex_project=vertex_ai_project, + vertex_location=vertex_ai_location, + vertex_credentials=vertex_credentials, + timeout=timeout, + max_retries=optional_params.max_retries, + create_file_data=_create_file_request, + ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus'] are supported.".format( @@ -362,7 +335,11 @@ def file_retrieve( _is_async = kwargs.pop("is_async", False) is True if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = _get_openai_credentials(optional_params) + openai_creds = get_openai_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + organization=optional_params.organization, + ) response = openai_files_instance.retrieve_file( file_id=file_id, _is_async=_is_async, @@ -373,7 +350,11 @@ def file_retrieve( organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - azure_creds = _get_azure_credentials(optional_params) + azure_creds = get_azure_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + api_version=optional_params.api_version, + ) response = azure_files_instance.retrieve_file( _is_async=_is_async, api_base=azure_creds.api_base, @@ -530,7 +511,11 @@ def file_delete( timeout = 600.0 _is_async = kwargs.pop("is_async", False) is True if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = _get_openai_credentials(optional_params) + openai_creds = get_openai_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + organization=optional_params.organization, + ) response = openai_files_instance.delete_file( file_id=file_id, _is_async=_is_async, @@ -541,7 +526,11 @@ def file_delete( organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - azure_creds = _get_azure_credentials(optional_params) + azure_creds = get_azure_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + api_version=optional_params.api_version, + ) response = azure_files_instance.delete_file( _is_async=_is_async, api_base=azure_creds.api_base, @@ -729,7 +718,11 @@ def file_list( ) return response elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = _get_openai_credentials(optional_params) + openai_creds = get_openai_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + organization=optional_params.organization, + ) response = openai_files_instance.list_files( purpose=purpose, _is_async=_is_async, @@ -740,7 +733,11 @@ def file_list( organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - azure_creds = _get_azure_credentials(optional_params) + azure_creds = get_azure_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + api_version=optional_params.api_version, + ) response = azure_files_instance.list_files( _is_async=_is_async, api_base=azure_creds.api_base, @@ -876,7 +873,11 @@ def file_content( return response if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = _get_openai_credentials(optional_params) + openai_creds = get_openai_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + organization=optional_params.organization, + ) response = openai_files_instance.file_content( _is_async=_is_async, file_content_request=_file_content_request, @@ -887,7 +888,11 @@ def file_content( organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - azure_creds = _get_azure_credentials(optional_params) + azure_creds = get_azure_credentials( + api_base=optional_params.api_base, + api_key=optional_params.api_key, + api_version=optional_params.api_version, + ) response = azure_files_instance.file_content( _is_async=_is_async, api_base=azure_creds.api_base, diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 25b218fca8c..7ed4306e299 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -1,6 +1,6 @@ import json import os -from typing import Any, Callable, Dict, Literal, Optional, Union, cast +from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Union, cast import httpx from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI @@ -789,3 +789,39 @@ class BaseAzureLLM(BaseOpenAILLM): return param_value return os.getenv(env_var_key) + +class AzureCredentials(NamedTuple): + api_base: Optional[str] + api_key: Optional[str] + api_version: Optional[str] + + +def get_azure_credentials( + api_base: Optional[str] = None, + api_key: Optional[str] = None, + api_version: Optional[str] = None, +) -> AzureCredentials: + """Resolve Azure credentials from params, litellm globals, and env vars.""" + resolved_api_base = ( + api_base + or litellm.api_base + or get_secret_str("AZURE_API_BASE") + ) + resolved_api_version = ( + api_version + or litellm.api_version + or get_secret_str("AZURE_API_VERSION") + ) + resolved_api_key = ( + api_key + or litellm.api_key + or litellm.azure_key + or get_secret_str("AZURE_OPENAI_API_KEY") + or get_secret_str("AZURE_API_KEY") + ) + return AzureCredentials( + api_base=resolved_api_base, + api_key=resolved_api_key, + api_version=resolved_api_version, + ) + diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 61f150f1c2e..069448eba07 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -5,8 +5,9 @@ Common helpers / utils across al OpenAI endpoints import hashlib import inspect import json +import os import ssl -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, NamedTuple, Optional, Tuple, Union import httpx import openai @@ -244,3 +245,39 @@ class BaseOpenAILLM: ) +class OpenAICredentials(NamedTuple): + api_base: Optional[str] + api_key: Optional[str] + organization: Optional[str] + + +def get_openai_credentials( + api_base: Optional[str] = None, + api_key: Optional[str] = None, + organization: Optional[str] = None, +) -> OpenAICredentials: + """Resolve OpenAI credentials from params, litellm globals, and env vars.""" + resolved_api_base = ( + api_base + or litellm.api_base + or os.getenv("OPENAI_BASE_URL") + or os.getenv("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + resolved_organization = ( + organization + or litellm.organization + or os.getenv("OPENAI_ORGANIZATION", None) + or None + ) + resolved_api_key = ( + api_key + or litellm.api_key + or litellm.openai_key + or os.getenv("OPENAI_API_KEY") + ) + return OpenAICredentials( + api_base=resolved_api_base, + api_key=resolved_api_key, + organization=resolved_organization, + ) From 5005773909e65a4ae1c87abcf86928c7052450e2 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 1 Dec 2025 11:32:55 -0300 Subject: [PATCH 079/147] fix(deps): relax python-multipart version constraint to >=0.0.22 The caret operator (^0.0.x) in zerover projects restricts to a single patch version. Changed to >= to allow future patch updates. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 577e51a0d22..395bbfccbd0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ orjson = {version = "^3.9.7", optional = true} apscheduler = {version = "^3.10.4", optional = true} fastapi-sso = { version = "^0.16.0", optional = true } PyJWT = { version = "^2.10.1", optional = true, python = ">=3.9" } -python-multipart = { version = "^0.0.22", optional = true, python = ">=3.10"} +python-multipart = { version = ">=0.0.20", optional = true} cryptography = {version = "*", optional = true} prisma = {version = "0.11.0", optional = true} azure-identity = {version = "^1.15.0", optional = true, python = ">=3.9"} From dad7805b42beac2e562253a6edd8df46cfb99436 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 2 Mar 2026 19:39:13 -0300 Subject: [PATCH 080/147] fix(deps): update python-multipart version to 0.0.22 in all files Align requirements.txt, CI workflow, liccheck, and license cache with the >=0.0.22 constraint already set in pyproject.toml. --- .github/workflows/test-litellm.yml | 2 +- requirements.txt | 3 ++- tests/code_coverage_tests/liccheck.ini | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index cf6928897be..3f8369df926 100644 --- a/.github/workflows/test-litellm.yml +++ b/.github/workflows/test-litellm.yml @@ -38,7 +38,7 @@ jobs: poetry run pip install "google-genai==1.22.0" poetry run pip install "google-cloud-aiplatform>=1.38" poetry run pip install "fastapi-offline==1.7.3" - poetry run pip install "python-multipart==0.0.22" + poetry run pip install "python-multipart>=0.0.20" poetry run pip install "openapi-core" - name: Setup litellm-enterprise as local package run: | diff --git a/requirements.txt b/requirements.txt index 69aac377d8a..3655c9b0270 100644 --- a/requirements.txt +++ b/requirements.txt @@ -41,8 +41,9 @@ polars==1.31.0 # for data processing apscheduler==3.10.4 # for resetting budget in background fastapi-sso==0.19.0 # admin UI, SSO pyjwt[crypto]==2.10.1 ; python_version >= "3.9" -python-multipart==0.0.22 # admin UI +python-multipart>=0.0.20 # admin UI jaraco.context>=6.1.0 +Pillow==11.0.0 azure-ai-contentsafety==1.0.0 # for azure content safety azure-identity==1.16.1 ; python_version >= "3.9" # for azure content safety azure-keyvault==4.2.0 # for azure KMS integration diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index 376d2859ffa..edb53ef15ad 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -114,7 +114,7 @@ apscheduler: >=3.10.4 # Unknown license fastapi-sso: >=0.16.0 # Unknown license filelock: >=3.20.0 # Unlicense (public domain) - https://unlicense.org / https://github.com/tox-dev/filelock pyjwt: >=2.9.0 # Unknown license -python-multipart: >=0.0.18 # Unknown license +python-multipart: >=0.0.22 # Unknown license pillow: >=11.0.0 # Unknown license azure-ai-contentsafety: >=1.0.0 # Unknown license azure-identity: >=1.16.1 # Unknown license From dc9f5a5cc41fd8b2c95d37e6ca1d501453ec820e Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 2 Mar 2026 20:25:22 -0300 Subject: [PATCH 081/147] fix(deps): update python-multipart to >=0.0.20 in CI and test configs --- poetry.lock | 17 +++++++++++++++-- tests/code_coverage_tests/liccheck.ini | 2 +- tests/code_coverage_tests/license_cache.json | 2 +- 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/poetry.lock b/poetry.lock index 3062c5fdaea..da511e2ebe5 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. [[package]] name = "a2a-sdk" @@ -5599,6 +5599,19 @@ files = [ [package.extras] cli = ["click (>=5.0)"] +[[package]] +name = "python-multipart" +version = "0.0.20" +description = "A streaming multipart parser for Python" +optional = true +python-versions = ">=3.8" +groups = ["main"] +markers = "python_version == \"3.9\" and extra == \"proxy\"" +files = [ + {file = "python_multipart-0.0.20-py3-none-any.whl", hash = "sha256:8a62d3a8335e06589fe01f2a3e178cdcc632f3fbe0d492ad9ee0ec35aab1f104"}, + {file = "python_multipart-0.0.20.tar.gz", hash = "sha256:8dd0cab45b8e23064ae09147625994d090fa46f5b0d1e13af944c331a7fa9d13"}, +] + [[package]] name = "python-multipart" version = "0.0.22" @@ -7980,4 +7993,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "5ae4b43dfe73be01d71f757227eb22245d18c06b5b4d5989b014500f400f1ee9" +content-hash = "70ec9abe5b06e7e81a2d76305cb950eea79692ae40321bac3285dc63fcbcf059" diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index edb53ef15ad..65ac01123d1 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -114,7 +114,7 @@ apscheduler: >=3.10.4 # Unknown license fastapi-sso: >=0.16.0 # Unknown license filelock: >=3.20.0 # Unlicense (public domain) - https://unlicense.org / https://github.com/tox-dev/filelock pyjwt: >=2.9.0 # Unknown license -python-multipart: >=0.0.22 # Unknown license +python-multipart: >=0.0.20 # Unknown license pillow: >=11.0.0 # Unknown license azure-ai-contentsafety: >=1.0.0 # Unknown license azure-identity: >=1.16.1 # Unknown license diff --git a/tests/code_coverage_tests/license_cache.json b/tests/code_coverage_tests/license_cache.json index a9c7fad2b14..e7b1157a240 100644 --- a/tests/code_coverage_tests/license_cache.json +++ b/tests/code_coverage_tests/license_cache.json @@ -20,7 +20,7 @@ "apscheduler:3.10.4": "MIT", "fastapi-sso:0.16.0": "MIT", "pyjwt:2.9.0": "MIT", - "python-multipart:0.0.22": "Apache-2.0", + "python-multipart:0.0.20": "Apache-2.0", "Pillow:11.0.0": "MIT-CMU", "azure-ai-contentsafety:1.0.0": "MIT License", "azure-identity:1.16.1": "MIT License", From 7d65df351fd852f7b544e126c763825296c8d7a5 Mon Sep 17 00:00:00 2001 From: Varad Khonde Date: Tue, 3 Mar 2026 23:40:41 +0530 Subject: [PATCH 082/147] feat(togetherai): add support for TogetherAI Qwen3.5-397B-A17B model --- litellm/model_prices_and_context_window_backup.json | 12 ++++++++++++ model_prices_and_context_window.json | 12 ++++++++++++ 2 files changed, 24 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5de764c5cec..3dfbc5cb219 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29736,6 +29736,18 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 412f99791a0..3a2d6d1d1e2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29958,6 +29958,18 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", From 224c61711948fd1508b59083542f022d1f37b99c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 10:12:08 -0800 Subject: [PATCH 083/147] Fix spend log cleanup: lock tracking, integer retention, skip log level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Only release distributed lock in finally if it was actually acquired; prevents spurious Redis release_lock calls on early returns - Treat bare integer maximum_spend_logs_retention_period as days (e.g. 3 → "3d") instead of silently failing with a ValueError - Elevate "Skipping cleanup" log from info to error so misconfigured retention settings are visible without verbose logging - Add tests for all three fixes Co-Authored-By: Claude Sonnet 4.6 --- .../db_transaction_queue/spend_log_cleanup.py | 13 ++-- .../proxy/test_spend_log_cleanup.py | 59 +++++++++++++++++++ 2 files changed, 68 insertions(+), 4 deletions(-) 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 8c59c79ff0a..6db7d6dd43c 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -49,7 +49,11 @@ class SpendLogCleanup: try: if isinstance(retention_setting, int): - retention_setting = str(retention_setting) + verbose_proxy_logger.warning( + f"maximum_spend_logs_retention_period is an integer ({retention_setting}); treating as days. " + "Use a string like '3d' to be explicit." + ) + retention_setting = f"{retention_setting}d" self.retention_seconds = duration_in_seconds(retention_setting) verbose_proxy_logger.info( f"Retention period set to {self.retention_seconds} seconds" @@ -112,11 +116,12 @@ class SpendLogCleanup: If pod_lock_manager is available, ensures only one pod runs cleanup. If no pod_lock_manager, runs cleanup without distributed locking. """ + lock_acquired = False try: verbose_proxy_logger.info(f"Cleanup job triggered at {datetime.now()}") if not self._should_delete_spend_logs(): - verbose_proxy_logger.info( + verbose_proxy_logger.error( "Skipping cleanup — invalid or missing retention setting." ) return @@ -155,8 +160,8 @@ class SpendLogCleanup: verbose_proxy_logger.error(f"Error during cleanup: {str(e)}") return # Return after error handling finally: - # Always release the lock if we have a pod lock manager - if self.pod_lock_manager and self.pod_lock_manager.redis_cache: + # Only release the lock if it was actually acquired + if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: await self.pod_lock_manager.release_lock( cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME ) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index c1fa3ad0c43..3a01437908d 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -233,6 +233,65 @@ async def test_cleanup_old_spend_logs_no_retention_period(): mock_prisma_client.db.execute_raw.assert_not_called() +@pytest.mark.asyncio +async def test_lock_not_released_when_not_acquired(): + """ + Lock release should be skipped when _should_delete_spend_logs returns False + before the lock is ever acquired. + """ + mock_prisma_client = MagicMock() + mock_prisma_client.db.execute_raw = AsyncMock() + + mock_redis_cache = MagicMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = mock_redis_cache + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + # No retention setting → _should_delete_spend_logs() returns False before lock is acquired + cleaner = SpendLogCleanup(general_settings={}) + cleaner.pod_lock_manager = mock_pod_lock_manager + + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + mock_pod_lock_manager.acquire_lock.assert_not_called() + mock_pod_lock_manager.release_lock.assert_not_called() + + +@pytest.mark.asyncio +async def test_integer_retention_treated_as_days(): + """ + An integer value for maximum_spend_logs_retention_period should be treated + as days (e.g., 3 → '3d' → 259200 seconds). + """ + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": 3} + ) + result = cleaner._should_delete_spend_logs() + assert result is True + assert cleaner.retention_seconds == 3 * 86400 # 3 days in seconds + + +def test_string_retention_still_works(): + """ + String values like '3d', '24h', '3600s' should continue to parse correctly. + """ + cases = [ + ("3d", 3 * 86400), + ("24h", 24 * 3600), + ("3600s", 3600), + ("2w", 2 * 604800), + ] + for setting, expected_seconds in cases: + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": setting} + ) + assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}" + assert cleaner.retention_seconds == expected_seconds, ( + f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + ) + + def test_cleanup_batch_size_env_var(monkeypatch): """Ensure batch size is configurable via environment variable""" import importlib From 47f0390b9b1320a586c9538936ae3a3b072af3b9 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 15:14:20 -0300 Subject: [PATCH 084/147] fix: remove duplicate Pillow==11.0.0 pin (12.1.1 already on line 8) --- requirements.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 3655c9b0270..aef0e1d271e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -43,7 +43,6 @@ fastapi-sso==0.19.0 # admin UI, SSO pyjwt[crypto]==2.10.1 ; python_version >= "3.9" python-multipart>=0.0.20 # admin UI jaraco.context>=6.1.0 -Pillow==11.0.0 azure-ai-contentsafety==1.0.0 # for azure content safety azure-identity==1.16.1 ; python_version >= "3.9" # for azure content safety azure-keyvault==4.2.0 # for azure KMS integration From fecb3016847d1c8faea966c21581242748801118 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Mar 2026 15:16:12 -0300 Subject: [PATCH 085/147] Update litellm/llms/openrouter/image_edit/transformation.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/llms/openrouter/image_edit/transformation.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index d06536367eb..19931f00245 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -133,12 +133,11 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): api_base: Optional[str], litellm_params: dict, ) -> str: - if api_base: - if not api_base.endswith("/chat/completions"): - api_base = api_base.rstrip("/") - return f"{api_base}/chat/completions" - return api_base - return "https://openrouter.ai/api/v1/chat/completions" + base_url = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" + base_url = base_url.rstrip("/") + if not base_url.endswith("/chat/completions"): + return f"{base_url}/chat/completions" + return base_url def transform_image_edit_request( self, From 4a88d854462be6cc1b8d7316a9df26ac1e65a487 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 15:20:51 -0300 Subject: [PATCH 086/147] test: add provider_url routing test for vertex_ai/gemini models Verifies that vertex_ai gemini models route to aiplatform.googleapis.com instead of generativelanguage.googleapis.com, preventing regressions if the branch ordering changes. --- .../helicone/test_helicone_gemini.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/tests/litellm/integrations/helicone/test_helicone_gemini.py b/tests/litellm/integrations/helicone/test_helicone_gemini.py index f42a7016131..67c4515c1e7 100644 --- a/tests/litellm/integrations/helicone/test_helicone_gemini.py +++ b/tests/litellm/integrations/helicone/test_helicone_gemini.py @@ -62,3 +62,74 @@ def test_helicone_vertex_ai_via_custom_llm_provider(): for model, custom_llm_provider in test_cases: is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith("vertex_ai/") assert is_vertex_ai, f"{model} with custom_llm_provider={custom_llm_provider} should be recognized as vertex_ai" + + +def test_helicone_vertex_gemini_gets_vertex_provider_url(): + """ + Test that vertex_ai/gemini-* models route to aiplatform.googleapis.com, + not generativelanguage.googleapis.com. + + This verifies the branch ordering fix: is_vertex_ai must be checked + before "gemini" in model, otherwise vertex gemini models get the wrong + provider_url. + """ + from unittest.mock import MagicMock, patch + + from litellm.integrations.helicone import HeliconeLogger + + logger = HeliconeLogger() + + captured = {} + + def mock_post(url, **kwargs): + captured["url"] = url + captured["data"] = kwargs.get("json", {}) + mock_resp = MagicMock() + mock_resp.status_code = 200 + return mock_resp + + test_cases = [ + # (model, custom_llm_provider, expected_provider_url) + ( + "vertex_ai/gemini-1.5-pro", + "", + "https://aiplatform.googleapis.com/v1", + ), + ( + "gemini-2.0-flash", + "vertex_ai", + "https://aiplatform.googleapis.com/v1", + ), + ( + "gemini-1.5-flash", + "", + "https://generativelanguage.googleapis.com/v1beta", + ), + ] + + for model, custom_llm_provider, expected_url in test_cases: + captured.clear() + mock_client = MagicMock() + mock_client.post = mock_post + with patch("litellm.module_level_client", mock_client): + logger.log_success( + model=model, + messages=[{"role": "user", "content": "test"}], + response_obj={"choices": [{"message": {"content": "hi"}}]}, + start_time=MagicMock(), + end_time=MagicMock(), + print_verbose=lambda *args, **kwargs: None, + kwargs={ + "litellm_params": { + "custom_llm_provider": custom_llm_provider, + "metadata": {}, + }, + }, + ) + + assert "data" in captured, f"No request captured for {model}" + actual_url = captured["data"]["providerRequest"]["url"] + assert actual_url == expected_url, ( + f"Model {model} (provider={custom_llm_provider!r}): " + f"expected provider_url={expected_url}, got {actual_url}" + ) From a1ba6c9fa643c04492d274b2a664d2de88ce1603 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 10:21:31 -0800 Subject: [PATCH 087/147] Fix log levels: info for unconfigured, warning for misconfigured Suppress noisy error log fired every cron tick when spend log cleanup is simply not configured. _should_delete_spend_logs already logs the specific reason at the right level (info for None, warning for invalid value), so the redundant blanket error log in cleanup_old_spend_logs is removed. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 6db7d6dd43c..a538e411b68 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -60,7 +60,7 @@ class SpendLogCleanup: ) return True except ValueError as e: - verbose_proxy_logger.error( + verbose_proxy_logger.warning( f"Invalid maximum_spend_logs_retention_period value: {retention_setting}, error: {str(e)}" ) return False @@ -121,9 +121,6 @@ class SpendLogCleanup: verbose_proxy_logger.info(f"Cleanup job triggered at {datetime.now()}") if not self._should_delete_spend_logs(): - verbose_proxy_logger.error( - "Skipping cleanup — invalid or missing retention setting." - ) return if self.retention_seconds is None: From a6e18ae31b53014534b7af51bc7b7ef17f870b46 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 15:36:45 -0300 Subject: [PATCH 088/147] =?UTF-8?q?fix:=20address=20review=20feedback=20?= =?UTF-8?q?=E2=80=94=20type=20OpenAICredentials.api=5Fbase=20as=20str,=20r?= =?UTF-8?q?emove=20dead=20vertex=5Fai=20block?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - OpenAICredentials.api_base is always non-None due to fallback default, so type it as str instead of Optional[str] - Remove unreachable elif vertex_ai block in create_file(); vertex_ai is already handled by ProviderConfigManager via provider_config path --- litellm/files/main.py | 26 -------------------------- litellm/llms/openai/common_utils.py | 2 +- 2 files changed, 1 insertion(+), 27 deletions(-) diff --git a/litellm/files/main.py b/litellm/files/main.py index f626f9b3466..1f4d15abb81 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -216,32 +216,6 @@ def create_file( create_file_data=_create_file_request, litellm_params=litellm_params_dict, ) - elif custom_llm_provider == "vertex_ai": - api_base = optional_params.api_base or "" - vertex_ai_project = ( - optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") - ) - vertex_ai_location = ( - optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str( - "VERTEXAI_CREDENTIALS" - ) - - response = vertex_ai_files_instance.create_file( - _is_async=_is_async, - api_base=api_base, - vertex_project=vertex_ai_project, - vertex_location=vertex_ai_location, - vertex_credentials=vertex_credentials, - timeout=timeout, - max_retries=optional_params.max_retries, - create_file_data=_create_file_request, - ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus'] are supported.".format( diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 069448eba07..b6b302782e8 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -246,7 +246,7 @@ class BaseOpenAILLM: class OpenAICredentials(NamedTuple): - api_base: Optional[str] + api_base: str api_key: Optional[str] organization: Optional[str] From 6f74d3881a4b6f459dd2f9e90c8a2f3a132b0602 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 15:42:45 -0300 Subject: [PATCH 089/147] fix(lint): remove unused imports in batch_utils.py time, httpx, uuid, and ModelResponse are no longer used after the vertex_ai batch cost calculation was rewritten upstream. --- litellm/batches/batch_utils.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 80351664dfe..a55e30ebeb9 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,14 +1,10 @@ import json -import time from typing import Any, List, Literal, Optional, Tuple -import httpx - import litellm from litellm._logging import verbose_logger -from litellm._uuid import uuid from litellm.types.llms.openai import Batch -from litellm.types.utils import CallTypes, ModelInfo, ModelResponse, Usage +from litellm.types.utils import CallTypes, ModelInfo, Usage from litellm.utils import token_counter From 43cec8c980abf699b3664a239393f408b9ca2f90 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Mar 2026 10:46:45 -0800 Subject: [PATCH 090/147] feat(batches): support output_expires_after passthrough --- litellm/batches/main.py | 5 ++ litellm/types/llms/openai.py | 1 + tests/test_litellm/proxy/test_batch_expiry.py | 74 +++++++++++++++++++ 3 files changed, 80 insertions(+) create mode 100644 tests/test_litellm/proxy/test_batch_expiry.py diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 9553d2c5246..e73f73d2f33 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -112,6 +112,7 @@ async def acreate_batch( metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, + output_expires_after: Optional[Dict[str, Any]] = None, **kwargs, ) -> LiteLLMBatch: """ @@ -133,6 +134,7 @@ async def acreate_batch( metadata, extra_headers, extra_body, + output_expires_after, **kwargs, ) @@ -160,6 +162,7 @@ def create_batch( metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, + output_expires_after: Optional[Dict[str, Any]] = None, **kwargs, ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: """ @@ -215,6 +218,8 @@ def create_batch( extra_headers=extra_headers, extra_body=extra_body, ) + if output_expires_after is not None: + _create_batch_request["output_expires_after"] = output_expires_after if model is not None: provider_config = ProviderConfigManager.get_provider_batches_config( model=model, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index d06d879dad1..c5d610e639b 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -424,6 +424,7 @@ class CreateBatchRequest(TypedDict, total=False): endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"] input_file_id: str metadata: Optional[Dict[str, str]] + output_expires_after: Optional[FileExpiresAfter] extra_headers: Optional[Dict[str, str]] extra_body: Optional[Dict[str, str]] timeout: Optional[float] diff --git a/tests/test_litellm/proxy/test_batch_expiry.py b/tests/test_litellm/proxy/test_batch_expiry.py new file mode 100644 index 00000000000..25b1631792e --- /dev/null +++ b/tests/test_litellm/proxy/test_batch_expiry.py @@ -0,0 +1,74 @@ +""" +Tests for batch output_expires_after passthrough and team-level expiry enforcement. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.types.llms.openai import CreateBatchRequest + + +class TestCreateBatchOutputExpiresAfterPassthrough: + """Verify output_expires_after flows through create_batch to the provider.""" + + def test_output_expires_after_included_in_request(self): + """When output_expires_after is provided, it reaches the openai batches instance.""" + captured = {} + + original_create = None + + def capturing_create(**kwargs): + captured.update(kwargs) + mock_response = MagicMock() + mock_response.id = "batch_123" + return mock_response + + with patch( + "litellm.batches.main.openai_batches_instance" + ) as mock_instance: + mock_instance.create_batch.side_effect = capturing_create + litellm.create_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-abc123", + output_expires_after={"anchor": "created_at", "seconds": 86400}, + custom_llm_provider="openai", + ) + + create_batch_data = captured["create_batch_data"] + assert create_batch_data["output_expires_after"] == { + "anchor": "created_at", + "seconds": 86400, + } + + def test_output_expires_after_absent_when_not_provided(self): + """Backward compat: output_expires_after not in request when omitted.""" + captured = {} + + def capturing_create(**kwargs): + captured.update(kwargs) + mock_response = MagicMock() + mock_response.id = "batch_123" + return mock_response + + with patch( + "litellm.batches.main.openai_batches_instance" + ) as mock_instance: + mock_instance.create_batch.side_effect = capturing_create + litellm.create_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-abc123", + custom_llm_provider="openai", + ) + + create_batch_data = captured["create_batch_data"] + assert "output_expires_after" not in create_batch_data From 4b1389501d3f1418077a63b1fac2151a37c8df89 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 15:58:23 -0300 Subject: [PATCH 091/147] fix: move _set_usage_and_cost outside try/except in OpenRouter image edit Prevents discarding already-parsed images when usage calculation fails. --- litellm/llms/openrouter/image_edit/transformation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index 19931f00245..ed5e6ae67d5 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -242,9 +242,6 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): ) ) - self._set_usage_and_cost(model_response, response_json, model) - return model_response - except Exception as e: raise OpenRouterException( message=f"Error transforming OpenRouter image edit response: {str(e)}", @@ -252,6 +249,9 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): headers={}, ) + self._set_usage_and_cost(model_response, response_json, model) + return model_response + def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: From 3d15bcdb115731cd0b4dd93f19be880e2128f1bb Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Mar 2026 10:58:31 -0800 Subject: [PATCH 092/147] feat(proxy): add team-level batch output expiry enforcement --- litellm/proxy/_types.py | 6 + litellm/proxy/batches_endpoints/endpoints.py | 9 + tests/test_litellm/proxy/test_batch_expiry.py | 188 +++++++++++++----- 3 files changed, 153 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index cbf683d226e..5122c64ea64 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1551,6 +1551,8 @@ class NewTeamRequest(TeamBase): ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None + enforced_batch_output_expires_after: Optional[dict] = None + enforced_file_expires_after: Optional[dict] = None model_config = ConfigDict(protected_namespaces=()) @@ -1606,6 +1608,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): model_rpm_limit: Optional[Dict[str, int]] = None model_tpm_limit: Optional[Dict[str, int]] = None allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None + enforced_batch_output_expires_after: Optional[dict] = None + enforced_file_expires_after: Optional[dict] = None router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None @@ -3783,6 +3787,8 @@ LiteLLM_ManagementEndpoint_MetadataFields = [ "temp_budget_increase", "temp_budget_expiry", "allowed_vector_store_indexes", + "enforced_batch_output_expires_after", + "enforced_file_expires_after", ] LiteLLM_ManagementEndpoint_MetadataFields_Premium = [ diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 1c9ba6cb248..60905243369 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -118,6 +118,15 @@ async def create_batch( # noqa: PLR0915 or "openai" ) _create_batch_data = LiteLLMBatchCreateRequest(**data) + + # Apply team-level batch output expiry enforcement + team_metadata = user_api_key_dict.team_metadata or {} + enforced_batch_expiry = team_metadata.get( + "enforced_batch_output_expires_after" + ) + if enforced_batch_expiry is not None: + _create_batch_data["output_expires_after"] = enforced_batch_expiry + input_file_id = _create_batch_data.get("input_file_id", None) unified_file_id: Union[str, Literal[False]] = False diff --git a/tests/test_litellm/proxy/test_batch_expiry.py b/tests/test_litellm/proxy/test_batch_expiry.py index 25b1631792e..1f54f190c63 100644 --- a/tests/test_litellm/proxy/test_batch_expiry.py +++ b/tests/test_litellm/proxy/test_batch_expiry.py @@ -4,7 +4,7 @@ Tests for batch output_expires_after passthrough and team-level expiry enforceme import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest @@ -13,62 +13,150 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm -from litellm.types.llms.openai import CreateBatchRequest +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.proxy_server import app +from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.types.utils import LiteLLMBatch + +from fastapi.testclient import TestClient + +client = TestClient(app) + +TEAM_EXPIRY = {"anchor": "created_at", "seconds": 3600} +CALLER_EXPIRY = {"anchor": "created_at", "seconds": 86400} -class TestCreateBatchOutputExpiresAfterPassthrough: - """Verify output_expires_after flows through create_batch to the provider.""" +@pytest.fixture +def llm_router() -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "test-key", + }, + "model_info": {"id": "gpt-3.5-turbo-id"}, + }, + ] + ) - def test_output_expires_after_included_in_request(self): - """When output_expires_after is provided, it reaches the openai batches instance.""" - captured = {} - original_create = None +def _setup_proxy(monkeypatch, llm_router: Router): + proxy_logging_obj = ProxyLogging( + user_api_key_cache=DualCache(default_in_memory_ttl=1) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ) - def capturing_create(**kwargs): - captured.update(kwargs) - mock_response = MagicMock() - mock_response.id = "batch_123" - return mock_response - with patch( - "litellm.batches.main.openai_batches_instance" - ) as mock_instance: - mock_instance.create_batch.side_effect = capturing_create - litellm.create_batch( - completion_window="24h", - endpoint="/v1/chat/completions", - input_file_id="file-abc123", - output_expires_after={"anchor": "created_at", "seconds": 86400}, - custom_llm_provider="openai", +def _make_batch_response() -> LiteLLMBatch: + return LiteLLMBatch( + id="batch_abc123", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id="file-abc123", + object="batch", + status="validating", + ) + + +def test_output_expires_after_passthrough(): + """output_expires_after flows through create_batch to the provider.""" + captured = {} + + def capturing_create(**kwargs): + captured.update(kwargs) + mock_response = MagicMock() + mock_response.id = "batch_123" + return mock_response + + with patch("litellm.batches.main.openai_batches_instance") as mock_instance: + mock_instance.create_batch.side_effect = capturing_create + litellm.create_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-abc123", + output_expires_after=CALLER_EXPIRY, + custom_llm_provider="openai", + ) + + assert captured["create_batch_data"]["output_expires_after"] == CALLER_EXPIRY + + +class TestBatchEndpointTeamOverride: + """Verify team-level enforced_batch_output_expires_after in the proxy endpoint.""" + + def _post_batch( + self, + monkeypatch, + llm_router: Router, + team_metadata: dict, + request_body: dict, + ) -> dict: + """POST /v1/batches with given team_metadata and body, return captured kwargs.""" + _setup_proxy(monkeypatch, llm_router) + + user_key = UserAPIKeyAuth( + api_key="test-key", + team_metadata=team_metadata, + ) + app.dependency_overrides[user_api_key_auth] = lambda: user_key + + captured_kwargs = {} + + async def mock_acreate_batch(**kwargs): + captured_kwargs.update(kwargs) + return _make_batch_response() + + monkeypatch.setattr(litellm, "acreate_batch", mock_acreate_batch) + + try: + response = client.post( + "/v1/batches", + json=request_body, + headers={"Authorization": "Bearer test-key"}, ) + assert response.status_code == 200 + finally: + app.dependency_overrides.clear() - create_batch_data = captured["create_batch_data"] - assert create_batch_data["output_expires_after"] == { - "anchor": "created_at", - "seconds": 86400, - } + return captured_kwargs - def test_output_expires_after_absent_when_not_provided(self): - """Backward compat: output_expires_after not in request when omitted.""" - captured = {} + def test_team_override_overrides_caller(self, monkeypatch, llm_router): + """Team enforcement wins over caller-provided value.""" + kwargs = self._post_batch( + monkeypatch, + llm_router, + team_metadata={ + "enforced_batch_output_expires_after": TEAM_EXPIRY, + }, + request_body={ + "input_file_id": "file-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "output_expires_after": CALLER_EXPIRY, + }, + ) + assert kwargs["output_expires_after"] == TEAM_EXPIRY - def capturing_create(**kwargs): - captured.update(kwargs) - mock_response = MagicMock() - mock_response.id = "batch_123" - return mock_response - - with patch( - "litellm.batches.main.openai_batches_instance" - ) as mock_instance: - mock_instance.create_batch.side_effect = capturing_create - litellm.create_batch( - completion_window="24h", - endpoint="/v1/chat/completions", - input_file_id="file-abc123", - custom_llm_provider="openai", - ) - - create_batch_data = captured["create_batch_data"] - assert "output_expires_after" not in create_batch_data + def test_no_team_setting_preserves_caller(self, monkeypatch, llm_router): + """No team setting = caller value passes through.""" + kwargs = self._post_batch( + monkeypatch, + llm_router, + team_metadata={}, + request_body={ + "input_file_id": "file-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "output_expires_after": CALLER_EXPIRY, + }, + ) + assert kwargs["output_expires_after"] == CALLER_EXPIRY From 08613b24cbbf0315249a305b743e43cc5e73411f Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Mar 2026 11:03:14 -0800 Subject: [PATCH 093/147] feat(proxy): add team-level file expiry enforcement --- .../openai_files_endpoints/files_endpoints.py | 11 +- .../test_files_endpoint.py | 135 ++++++++++++++++++ 2 files changed, 145 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index ec6e9733344..a641fb1c69f 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -454,8 +454,17 @@ async def create_file( # noqa: PLR0915 model=router_model, llm_router=llm_router ) + # Apply team-level file expiry enforcement + team_metadata = user_api_key_dict.team_metadata or {} + enforced_file_expiry = team_metadata.get("enforced_file_expires_after") + if enforced_file_expiry is not None: + expires_after = FileExpiresAfter( + anchor=enforced_file_expiry["anchor"], + seconds=enforced_file_expiry["seconds"], + ) + _create_file_request = CreateFileRequest( - file=file_data, + file=file_data, purpose=cast(CREATE_FILE_REQUESTS_PURPOSE, purpose), expires_after=expires_after, **data diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index fb063ef8ee7..0239c39e67f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1168,3 +1168,138 @@ def test_create_file_with_deep_nested_litellm_metadata( assert captured_litellm_metadata["config"]["database"]["port"] == "5432" assert "cache" in captured_litellm_metadata["config"] assert captured_litellm_metadata["config"]["cache"]["enabled"] == "true" + + +# --------------------------------------------------------------------------- +# Team-level enforced_file_expires_after tests +# --------------------------------------------------------------------------- + + +def _make_capturing_managed_files(): + """Create a DummyManagedFiles that captures the expires_after from the request.""" + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + + captured = {} + + class CapturingManagedFiles(BaseFileEndpoints): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): + if isinstance(create_file_request, dict): + captured["expires_after"] = create_file_request.get("expires_after") + else: + captured["expires_after"] = getattr( + create_file_request, "expires_after", None + ) + return OpenAIFileObject( + id="file-abc123", + object="file", + bytes=100, + created_at=1234567890, + filename="mydata.jsonl", + purpose="batch", + status="uploaded", + ) + + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): + raise NotImplementedError + + async def afile_list(self, purpose, litellm_parent_otel_span): + raise NotImplementedError + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): + raise NotImplementedError + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): + raise NotImplementedError + + return CapturingManagedFiles(), captured + + +def _post_file_with_team_metadata( + monkeypatch, + llm_router: Router, + team_metadata: dict, + form_data: dict, +): + """POST /v1/files with given team_metadata, return captured expires_after.""" + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + proxy_logging_obj = ProxyLogging( + user_api_key_cache=DualCache(default_in_memory_ttl=1) + ) + dummy, captured = _make_capturing_managed_files() + proxy_logging_obj.proxy_hook_mapping["managed_files"] = dummy + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ) + + user_key = UserAPIKeyAuth(api_key="test-key", team_metadata=team_metadata) + app.dependency_overrides[user_api_key_auth] = lambda: user_key + + test_file = ("mydata.jsonl", b'{"prompt": "Hello"}', "application/json") + try: + response = client.post( + "/v1/files", + files={"file": test_file}, + data=form_data, + headers={"Authorization": "Bearer test-key"}, + ) + assert response.status_code == 200 + finally: + app.dependency_overrides.clear() + + return captured["expires_after"] + + +def test_file_team_override_overrides_caller( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + """Team enforced_file_expires_after wins over caller-provided value.""" + expires_after = _post_file_with_team_metadata( + monkeypatch, + llm_router, + team_metadata={ + "enforced_file_expires_after": { + "anchor": "created_at", + "seconds": 3600, + } + }, + form_data={ + "purpose": "batch", + "target_model_names": "gpt-3.5-turbo", + "expires_after[anchor]": "created_at", + "expires_after[seconds]": "86400", + }, + ) + assert expires_after["anchor"] == "created_at" + assert expires_after["seconds"] == 3600 + + +def test_file_no_team_setting_preserves_caller( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + """No team setting = caller-provided expires_after passes through.""" + expires_after = _post_file_with_team_metadata( + monkeypatch, + llm_router, + team_metadata={}, + form_data={ + "purpose": "batch", + "target_model_names": "gpt-3.5-turbo", + "expires_after[anchor]": "created_at", + "expires_after[seconds]": "86400", + }, + ) + assert expires_after["anchor"] == "created_at" + assert expires_after["seconds"] == 86400 From d6614191096d2a9e06e7d9a36281ce71ad4170ab Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 4 Mar 2026 01:50:28 +0530 Subject: [PATCH 094/147] fix: support list of modes in Mode.default for tag-based guardrails --- .../integrations/custom_guardrail.py | 37 ++++- litellm/integrations/custom_guardrail.py | 4 +- .../integrations/test_custom_guardrail.py | 130 +++++++++++++++++- 3 files changed, 158 insertions(+), 13 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/custom_guardrail.py b/enterprise/litellm_enterprise/integrations/custom_guardrail.py index b165d788f35..8ed3bfcac4c 100644 --- a/enterprise/litellm_enterprise/integrations/custom_guardrail.py +++ b/enterprise/litellm_enterprise/integrations/custom_guardrail.py @@ -10,10 +10,15 @@ class EnterpriseCustomGuardrailHelper: event_hook: Optional[ Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode] ], + event_type: Optional[GuardrailEventHooks] = None, ) -> Optional[bool]: """ - Assumes check for event match is done in `should_run_guardrail` - Returns True if the guardrail should be run by tag + Returns True if the guardrail should be run for this request and event_type. + + Logic: + - If a request tag matches a Mode tag key, only run if event_type matches + the tag's value (the mode for that tag). + - If no request tag matches, fall back to default mode(s). """ from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, @@ -36,11 +41,29 @@ class EnterpriseCustomGuardrailHelper: proxy_server_request=proxy_server_request, ) - if request_tags and any(tag in event_hook.tags for tag in request_tags): - return True - elif event_hook.default and any( - tag in event_hook.default for tag in request_tags - ): + # Check if any request tag matches a Mode tag key + matched_mode = None + if request_tags: + for tag in request_tags: + if tag in event_hook.tags: + matched_mode = event_hook.tags[tag] + break + + if matched_mode is not None: + # Tag matched: only run if event_type matches the tag's mode value + if event_type is not None: + return event_type.value == matched_mode return True + # No tag matched: fall back to default mode(s) + if event_hook.default is not None: + if event_type is not None: + default_list = ( + event_hook.default + if isinstance(event_hook.default, list) + else [event_hook.default] + ) + return event_type.value in default_list + return False + return False diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 3fd179bb411..269797b9873 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -420,7 +420,7 @@ class CustomGuardrail(CustomLogger): "Setting tag-based guardrails is only available in litellm-enterprise. You must be a premium user to use this feature." ) result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag( - data, self.event_hook + data, self.event_hook, event_type ) if result is not None: return result @@ -447,7 +447,7 @@ class CustomGuardrail(CustomLogger): "Setting tag-based guardrails is only available in litellm-enterprise. You must be a premium user to use this feature." ) result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag( - data, self.event_hook + data, self.event_hook, event_type ) if result is not None: return result diff --git a/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py b/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py index 6feaca6f0b7..f4e06f9f317 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py @@ -1,9 +1,5 @@ -import datetime -import json import os import sys -import unittest -from unittest.mock import ANY, MagicMock, patch sys.path.insert( 0, os.path.abspath("../..") @@ -12,6 +8,132 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks, Mode +def test_custom_guardrail_with_mode_default_list(monkeypatch): + """Test Mode with default as a list of modes (e.g. default: ["pre_call", "post_call"])""" + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cg = CustomGuardrail( + guardrail_name="test_guardrail", + supported_event_hooks=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.logging_only, + ], + event_hook=Mode( + tags={"test_tag": "logging_only"}, + default=["pre_call", "post_call"], + ), + default_on=True, + ) + + # No tag match → default fires for pre_call + assert ( + cg.should_run_guardrail( + data={"messages": [{"role": "user", "content": "test"}]}, + event_type=GuardrailEventHooks.pre_call, + ) + is True + ) + + # No tag match → default fires for post_call + assert ( + cg.should_run_guardrail( + data={"messages": [{"role": "user", "content": "test"}]}, + event_type=GuardrailEventHooks.post_call, + ) + is True + ) + + # No tag match → logging_only NOT in default list, should not fire + assert ( + cg.should_run_guardrail( + data={"messages": [{"role": "user", "content": "test"}]}, + event_type=GuardrailEventHooks.logging_only, + ) + is False + ) + + # Tag matches → only logging_only should fire + assert ( + cg.should_run_guardrail( + data={ + "messages": [{"role": "user", "content": "test"}], + "litellm_metadata": {"tags": ["test_tag"]}, + }, + event_type=GuardrailEventHooks.logging_only, + ) + is True + ) + + # Tag matches → pre_call should NOT fire (tag says logging_only) + assert ( + cg.should_run_guardrail( + data={ + "messages": [{"role": "user", "content": "test"}], + "litellm_metadata": {"tags": ["test_tag"]}, + }, + event_type=GuardrailEventHooks.pre_call, + ) + is False + ) + + # Tag matches → post_call should NOT fire (tag says logging_only) + assert ( + cg.should_run_guardrail( + data={ + "messages": [{"role": "user", "content": "test"}], + "litellm_metadata": {"tags": ["test_tag"]}, + }, + event_type=GuardrailEventHooks.post_call, + ) + is False + ) + + +def test_custom_guardrail_with_mode_no_default(monkeypatch): + """Test Mode with no default — guardrail only fires when tag matches""" + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cg = CustomGuardrail( + guardrail_name="test_guardrail", + supported_event_hooks=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.logging_only, + ], + event_hook=Mode( + tags={"test_tag": "logging_only"}, + ), + default_on=True, + ) + + # No tag, no default → nothing fires + assert ( + cg.should_run_guardrail( + data={"messages": [{"role": "user", "content": "test"}]}, + event_type=GuardrailEventHooks.pre_call, + ) + is False + ) + + assert ( + cg.should_run_guardrail( + data={"messages": [{"role": "user", "content": "test"}]}, + event_type=GuardrailEventHooks.logging_only, + ) + is False + ) + + # Tag matches → only logging_only fires + assert ( + cg.should_run_guardrail( + data={ + "messages": [{"role": "user", "content": "test"}], + "litellm_metadata": {"tags": ["test_tag"]}, + }, + event_type=GuardrailEventHooks.logging_only, + ) + is True + ) + + def test_custom_guardrail_with_mode(monkeypatch): monkeypatch.setattr( "litellm.proxy.proxy_server.premium_user", True From 4edb1e00c60bf620ec6a179e29224298b000e081 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 12:36:26 -0800 Subject: [PATCH 095/147] [Test] UI - Add unit tests for project hooks Co-Authored-By: Claude Sonnet 4.6 --- .../hooks/projects/useCreateProject.test.ts | 111 ++++++++++++++ .../hooks/projects/useDeleteProject.test.ts | 88 +++++++++++ .../hooks/projects/useProjectDetails.test.ts | 144 ++++++++++++++++++ .../hooks/projects/useProjects.test.ts | 124 +++++++++++++++ .../hooks/projects/useUpdateProject.test.ts | 116 ++++++++++++++ 5 files changed, 583 insertions(+) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts new file mode 100644 index 00000000000..64d950d59ee --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts @@ -0,0 +1,111 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useCreateProject, ProjectCreateParams } from "./useCreateProject"; +import { projectKeys, ProjectResponse } from "./useProjects"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => ""), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), + deriveErrorMessage: vi.fn((data: any) => data?.error || "Error"), + handleError: vi.fn(), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +const mockProject: ProjectResponse = { + project_id: "proj-1", + project_alias: "Test Project", + description: "A test project", + team_id: "team-1", + budget_id: null, + metadata: null, + models: ["gpt-4"], + spend: 25.0, + model_spend: null, + model_rpm_limit: null, + model_tpm_limit: null, + blocked: false, + object_permission_id: null, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-02T00:00:00Z", + updated_by: "user-1", + litellm_budget_table: null, +}; + +function makeWrapper(queryClient: QueryClient) { + return ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); +} + +describe("useCreateProject", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + vi.clearAllMocks(); + global.fetch = vi.fn(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Admin" }); + }); + + it("should render", () => { + const { result } = renderHook(() => useCreateProject(), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current.mutate).toBeDefined(); + }); + + it("should POST to /project/new and return the created project", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProject }); + const { result } = renderHook(() => useCreateProject(), { + wrapper: makeWrapper(queryClient), + }); + const params: ProjectCreateParams = { team_id: "team-1", project_alias: "New Project" }; + const data = await result.current.mutateAsync(params); + expect(data).toEqual(mockProject); + const [url, init] = (global.fetch as any).mock.calls[0]; + expect(url).toContain("/project/new"); + expect(init.method).toBe("POST"); + expect(JSON.parse(init.body)).toMatchObject(params); + }); + + it("should invalidate project queries on success", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProject }); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + const { result } = renderHook(() => useCreateProject(), { + wrapper: makeWrapper(queryClient), + }); + await result.current.mutateAsync({ team_id: "team-1" }); + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: projectKeys.all }); + }); + + it("should set isError when the request fails", async () => { + (global.fetch as any).mockResolvedValue({ + ok: false, + json: async () => ({ error: "Server error" }), + }); + const { result } = renderHook(() => useCreateProject(), { + wrapper: makeWrapper(queryClient), + }); + result.current.mutateAsync({ team_id: "team-1" }).catch(() => {}); + await waitFor(() => expect(result.current.isError).toBe(true)); + }); + + it("should throw when accessToken is missing", async () => { + mockUseAuthorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { result } = renderHook(() => useCreateProject(), { + wrapper: makeWrapper(queryClient), + }); + await expect(result.current.mutateAsync({ team_id: "team-1" })).rejects.toThrow( + "Access token is required" + ); + expect(global.fetch).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.test.ts new file mode 100644 index 00000000000..85a9f3e0b10 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useDeleteProject.test.ts @@ -0,0 +1,88 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useDeleteProject } from "./useDeleteProject"; +import { projectKeys } from "./useProjects"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => ""), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), + deriveErrorMessage: vi.fn((data: any) => data?.error || "Error"), + handleError: vi.fn(), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +function makeWrapper(queryClient: QueryClient) { + return ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); +} + +describe("useDeleteProject", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + vi.clearAllMocks(); + global.fetch = vi.fn(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Admin" }); + }); + + it("should render", () => { + const { result } = renderHook(() => useDeleteProject(), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current.mutate).toBeDefined(); + }); + + it("should send DELETE to /project/delete with the given project IDs", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => ({}) }); + const { result } = renderHook(() => useDeleteProject(), { + wrapper: makeWrapper(queryClient), + }); + await result.current.mutateAsync(["proj-1", "proj-2"]); + const [url, init] = (global.fetch as any).mock.calls[0]; + expect(url).toContain("/project/delete"); + expect(init.method).toBe("DELETE"); + expect(JSON.parse(init.body)).toEqual({ project_ids: ["proj-1", "proj-2"] }); + }); + + it("should invalidate project queries on success", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => ({}) }); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + const { result } = renderHook(() => useDeleteProject(), { + wrapper: makeWrapper(queryClient), + }); + await result.current.mutateAsync(["proj-1"]); + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: projectKeys.all }); + }); + + it("should set isError when the request fails", async () => { + (global.fetch as any).mockResolvedValue({ + ok: false, + json: async () => ({ error: "Not found" }), + }); + const { result } = renderHook(() => useDeleteProject(), { + wrapper: makeWrapper(queryClient), + }); + result.current.mutateAsync(["proj-1"]).catch(() => {}); + await waitFor(() => expect(result.current.isError).toBe(true)); + }); + + it("should throw when accessToken is missing", async () => { + mockUseAuthorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { result } = renderHook(() => useDeleteProject(), { + wrapper: makeWrapper(queryClient), + }); + await expect(result.current.mutateAsync(["proj-1"])).rejects.toThrow( + "Access token is required" + ); + expect(global.fetch).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.test.ts new file mode 100644 index 00000000000..426abfe9bb6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.test.ts @@ -0,0 +1,144 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useProjectDetails } from "./useProjectDetails"; +import { projectKeys, ProjectResponse } from "./useProjects"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => ""), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), + deriveErrorMessage: vi.fn((data: any) => data?.error || "Error"), + handleError: vi.fn(), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +const mockProject: ProjectResponse = { + project_id: "proj-1", + project_alias: "Test Project", + description: "A test project", + team_id: "team-1", + budget_id: null, + metadata: null, + models: ["gpt-4"], + spend: 25.0, + model_spend: null, + model_rpm_limit: null, + model_tpm_limit: null, + blocked: false, + object_permission_id: null, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-02T00:00:00Z", + updated_by: "user-1", + litellm_budget_table: null, +}; + +const mockProjects: ProjectResponse[] = [ + mockProject, + { ...mockProject, project_id: "proj-2", project_alias: "Test Project 2" }, +]; + +function makeWrapper(queryClient: QueryClient) { + return ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); +} + +describe("useProjectDetails", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + vi.clearAllMocks(); + global.fetch = vi.fn(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Admin" }); + }); + + it("should render", () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProject }); + const { result } = renderHook(() => useProjectDetails("proj-1"), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current).toBeDefined(); + }); + + it("should return project details when the request succeeds", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProject }); + const { result } = renderHook(() => useProjectDetails("proj-1"), { + wrapper: makeWrapper(queryClient), + }); + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(result.current.data).toEqual(mockProject); + }); + + it("should call /project/info with the projectId encoded as a query param", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProject }); + renderHook(() => useProjectDetails("proj-1"), { wrapper: makeWrapper(queryClient) }); + await waitFor(() => expect(global.fetch).toHaveBeenCalled()); + const [url] = (global.fetch as any).mock.calls[0]; + expect(url).toContain("/project/info"); + expect(url).toContain("project_id=proj-1"); + }); + + it("should not fetch when projectId is missing", () => { + const { result } = renderHook(() => useProjectDetails(undefined), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current.isFetched).toBe(false); + expect(global.fetch).not.toHaveBeenCalled(); + }); + + it("should not fetch when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { result } = renderHook(() => useProjectDetails("proj-1"), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current.isFetched).toBe(false); + expect(global.fetch).not.toHaveBeenCalled(); + }); + + it("should not fetch when userRole is not an admin role", () => { + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Internal User" }); + const { result } = renderHook(() => useProjectDetails("proj-1"), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current.isFetched).toBe(false); + expect(global.fetch).not.toHaveBeenCalled(); + }); + + it("should seed initialData from the projects list cache", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProject }); + queryClient.setQueryData(projectKeys.list({}), mockProjects); + const { result } = renderHook(() => useProjectDetails("proj-1"), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current.data).toEqual(mockProject); + expect(result.current.isLoading).toBe(false); + await waitFor(() => expect(result.current.isFetching).toBe(false)); + }); + + it("should return undefined initialData when projectId is not in the cache", () => { + queryClient.setQueryData(projectKeys.list({}), mockProjects); + const { result } = renderHook(() => useProjectDetails("non-existent"), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current.data).toBeUndefined(); + }); + + it("should set isError when the request fails", async () => { + (global.fetch as any).mockResolvedValue({ + ok: false, + json: async () => ({ error: "Not found" }), + }); + const { result } = renderHook(() => useProjectDetails("proj-1"), { + wrapper: makeWrapper(queryClient), + }); + await waitFor(() => expect(result.current.isError).toBe(true)); + }); +}); 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 new file mode 100644 index 00000000000..13b9107bdc1 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts @@ -0,0 +1,124 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useProjects, ProjectResponse } from "./useProjects"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => ""), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), + deriveErrorMessage: vi.fn((data: any) => data?.error || "Error"), + handleError: vi.fn(), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +const mockProjects: ProjectResponse[] = [ + { + project_id: "proj-1", + project_alias: "Test Project", + description: "A test project", + team_id: "team-1", + budget_id: null, + metadata: null, + models: ["gpt-4"], + spend: 25.0, + model_spend: null, + model_rpm_limit: null, + model_tpm_limit: null, + blocked: false, + object_permission_id: null, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-02T00:00:00Z", + updated_by: "user-1", + litellm_budget_table: null, + }, + { + project_id: "proj-2", + project_alias: "Test Project 2", + description: null, + team_id: "team-1", + budget_id: null, + metadata: null, + models: [], + spend: 0, + model_spend: null, + model_rpm_limit: null, + model_tpm_limit: null, + blocked: false, + object_permission_id: null, + created_at: "2024-01-03T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-03T00:00:00Z", + updated_by: "user-1", + litellm_budget_table: null, + }, +]; + +function makeWrapper(queryClient: QueryClient) { + return ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); +} + +describe("useProjects", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + vi.clearAllMocks(); + global.fetch = vi.fn(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Admin" }); + }); + + it("should render", () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProjects }); + const { result } = renderHook(() => useProjects(), { wrapper: makeWrapper(queryClient) }); + expect(result.current).toBeDefined(); + }); + + it("should return projects when the request succeeds", async () => { + (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(result.current.data).toEqual(mockProjects); + }); + + it("should call GET /project/list with the auth header", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProjects }); + renderHook(() => useProjects(), { wrapper: makeWrapper(queryClient) }); + await waitFor(() => expect(global.fetch).toHaveBeenCalled()); + const [url, init] = (global.fetch as any).mock.calls[0]; + expect(url).toContain("/project/list"); + expect(init.headers["Authorization"]).toBe("Bearer test-token"); + }); + + it("should set isError when the request fails", async () => { + (global.fetch as any).mockResolvedValue({ + ok: false, + json: async () => ({ error: "Not authorized" }), + }); + const { result } = renderHook(() => useProjects(), { wrapper: makeWrapper(queryClient) }); + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(result.current.data).toBeUndefined(); + }); + + it("should not fetch when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { result } = renderHook(() => useProjects(), { wrapper: makeWrapper(queryClient) }); + expect(result.current.isFetched).toBe(false); + expect(global.fetch).not.toHaveBeenCalled(); + }); + + it("should not fetch when userRole is not an admin role", () => { + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Internal 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/useUpdateProject.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts new file mode 100644 index 00000000000..31d1a5fb352 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts @@ -0,0 +1,116 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useUpdateProject } from "./useUpdateProject"; +import { projectKeys, ProjectResponse } from "./useProjects"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => ""), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), + deriveErrorMessage: vi.fn((data: any) => data?.error || "Error"), + handleError: vi.fn(), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +const mockProject: ProjectResponse = { + project_id: "proj-1", + project_alias: "Test Project", + description: "A test project", + team_id: "team-1", + budget_id: null, + metadata: null, + models: ["gpt-4"], + spend: 25.0, + model_spend: null, + model_rpm_limit: null, + model_tpm_limit: null, + blocked: false, + object_permission_id: null, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-02T00:00:00Z", + updated_by: "user-1", + litellm_budget_table: null, +}; + +function makeWrapper(queryClient: QueryClient) { + return ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); +} + +describe("useUpdateProject", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + vi.clearAllMocks(); + global.fetch = vi.fn(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Admin" }); + }); + + it("should render", () => { + const { result } = renderHook(() => useUpdateProject(), { + wrapper: makeWrapper(queryClient), + }); + expect(result.current.mutate).toBeDefined(); + }); + + it("should POST to /project/update and return the updated project", async () => { + const updated = { ...mockProject, project_alias: "Updated Name" }; + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => updated }); + const { result } = renderHook(() => useUpdateProject(), { + wrapper: makeWrapper(queryClient), + }); + const data = await result.current.mutateAsync({ + projectId: "proj-1", + params: { project_alias: "Updated Name" }, + }); + expect(data).toEqual(updated); + const [url, init] = (global.fetch as any).mock.calls[0]; + expect(url).toContain("/project/update"); + expect(JSON.parse(init.body)).toMatchObject({ + project_id: "proj-1", + project_alias: "Updated Name", + }); + }); + + it("should invalidate project queries on success", async () => { + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProject }); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + const { result } = renderHook(() => useUpdateProject(), { + wrapper: makeWrapper(queryClient), + }); + await result.current.mutateAsync({ projectId: "proj-1", params: {} }); + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: projectKeys.all }); + }); + + it("should set isError when the request fails", async () => { + (global.fetch as any).mockResolvedValue({ + ok: false, + json: async () => ({ error: "Server error" }), + }); + const { result } = renderHook(() => useUpdateProject(), { + wrapper: makeWrapper(queryClient), + }); + result.current.mutateAsync({ projectId: "proj-1", params: {} }).catch(() => {}); + await waitFor(() => expect(result.current.isError).toBe(true)); + }); + + it("should throw when accessToken is missing", async () => { + mockUseAuthorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { result } = renderHook(() => useUpdateProject(), { + wrapper: makeWrapper(queryClient), + }); + await expect( + result.current.mutateAsync({ projectId: "proj-1", params: {} }) + ).rejects.toThrow("Access token is required"); + expect(global.fetch).not.toHaveBeenCalled(); + }); +}); From 279b4f16cb3c5d783294928f1018bf05320dfb50 Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Tue, 3 Mar 2026 17:57:32 -0300 Subject: [PATCH 096/147] Fix mypy override errors in count_tokens signatures Replace **kwargs with explicit tools and system parameters to match the BaseTokenCounter.count_tokens abstract method signature. Co-Authored-By: Claude Opus 4.6 --- litellm/llms/gemini/common_utils.py | 3 ++- litellm/llms/vertex_ai/common_utils.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index f99548c2c45..17b9c78123f 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -166,7 +166,8 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", - **kwargs, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Optional[TokenCountResponse]: import copy diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 791878c9700..3c5cbb65437 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1030,7 +1030,8 @@ class VertexAITokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", - **kwargs, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Optional[TokenCountResponse]: import copy From 6b4bc99202c1fad64920c260f437ca52116966a5 Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Tue, 3 Mar 2026 18:12:49 -0300 Subject: [PATCH 097/147] Fix Anthropic streaming sync __next__ and Azure GPT-5.1 logprobs Two independent fixes for pre-existing test failures on main: 1. Anthropic streaming: The sync __next__ method used a simple holding_chunk pattern that lost chunks when multiple events needed to be returned. Refactored to use the same chunk_queue approach as the async __anext__ method. Also fixed tests that used ModelResponse (which defaults finish_reason to 'stop') instead of ModelResponseStream. 2. Azure GPT-5.1 logprobs: The base OpenAI class includes logprobs for gpt-5.1+ models, but Azure hasn't verified support for gpt-5.1. Added explicit removal of logprobs/top_logprobs for gpt-5.1 (non-5.2) models in the Azure config. Co-Authored-By: Claude Opus 4.6 --- .../adapters/streaming_iterator.py | 135 ++++++++++-------- .../llms/azure/chat/gpt_5_transformation.py | 8 +- .../test_content_after_stop_reason.py | 14 +- .../messages/test_parallel_tool_calls.py | 26 ++-- .../messages/test_sse_wrapper.py | 17 +-- 5 files changed, 105 insertions(+), 95 deletions(-) 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 de634ff9ecf..cdf8ac5ca82 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -80,38 +80,40 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): from .transformation import LiteLLMAnthropicMessagesAdapter try: + # Always return queued chunks first + if self.chunk_queue: + return self.chunk_queue.popleft() + + # Queue initial chunks if not sent yet if self.sent_first_chunk is False: self.sent_first_chunk = True - return { - "type": "message_start", - "message": { - "id": "msg_{}".format(uuid.uuid4()), - "type": "message", - "role": "assistant", - "content": [], - "model": self.model, - "stop_reason": None, - "stop_sequence": None, - "usage": self._create_initial_usage_delta(), - }, - } + self.chunk_queue.append( + { + "type": "message_start", + "message": { + "id": "msg_{}".format(uuid.uuid4()), + "type": "message", + "role": "assistant", + "content": [], + "model": self.model, + "stop_reason": None, + "stop_sequence": None, + "usage": self._create_initial_usage_delta(), + }, + } + ) + return self.chunk_queue.popleft() + if self.sent_content_block_start is False: self.sent_content_block_start = True - return { - "type": "content_block_start", - "index": self.current_content_block_index, - "content_block": {"type": "text", "text": ""}, - } - - # Handle pending new content block start - if self.pending_new_content_block: - self.pending_new_content_block = False - self.sent_content_block_finish = False # Reset for new block - return { - "type": "content_block_start", - "index": self.current_content_block_index, - "content_block": self.current_content_block_start, - } + self.chunk_queue.append( + { + "type": "content_block_start", + "index": self.current_content_block_index, + "content_block": {"type": "text", "text": ""}, + } + ) + return self.chunk_queue.popleft() for chunk in self.completion_stream: if chunk == "None" or chunk is None: @@ -126,45 +128,65 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): current_content_block_index=self.current_content_block_index, ) - # Check if we need to start a new content block - # This is where you'd add your logic to detect when a new content block should start - # For example, if the chunk indicates a tool call or different content type - if should_start_new_block and not self.sent_content_block_finish: - # End current content block and prepare for new one - self.holding_chunk = processed_chunk - self.sent_content_block_finish = True - self.pending_new_content_block = True - return { - "type": "content_block_stop", - "index": max(self.current_content_block_index - 1, 0), - } + # Queue the sequence: content_block_stop -> content_block_start + # The trigger chunk itself is not emitted as a delta since the + # content_block_start already carries the relevant information. + self.chunk_queue.append( + { + "type": "content_block_stop", + "index": max(self.current_content_block_index - 1, 0), + } + ) + self.chunk_queue.append( + { + "type": "content_block_start", + "index": self.current_content_block_index, + "content_block": self.current_content_block_start, + } + ) + self.sent_content_block_finish = False + return self.chunk_queue.popleft() if ( processed_chunk["type"] == "message_delta" and self.sent_content_block_finish is False ): - self.holding_chunk = processed_chunk + # Queue both the content_block_stop and the message_delta + self.chunk_queue.append( + { + "type": "content_block_stop", + "index": self.current_content_block_index, + } + ) self.sent_content_block_finish = True - return { - "type": "content_block_stop", - "index": self.current_content_block_index, - } + self.chunk_queue.append(processed_chunk) + return self.chunk_queue.popleft() elif self.holding_chunk is not None: - return_chunk = self.holding_chunk - self.holding_chunk = processed_chunk - return return_chunk + self.chunk_queue.append(self.holding_chunk) + self.chunk_queue.append(processed_chunk) + self.holding_chunk = None + return self.chunk_queue.popleft() else: - return processed_chunk + self.chunk_queue.append(processed_chunk) + return self.chunk_queue.popleft() + + # Handle any remaining held chunks after stream ends if self.holding_chunk is not None: - return_chunk = self.holding_chunk + self.chunk_queue.append(self.holding_chunk) self.holding_chunk = None - return return_chunk - if self.sent_last_message is False: + + if not self.sent_last_message: self.sent_last_message = True - return {"type": "message_stop"} + self.chunk_queue.append({"type": "message_stop"}) + + if self.chunk_queue: + return self.chunk_queue.popleft() + raise StopIteration except StopIteration: + if self.chunk_queue: + return self.chunk_queue.popleft() if self.sent_last_message is False: self.sent_last_message = True return {"type": "message_stop"} @@ -265,7 +287,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if not self.queued_usage_chunk: if should_start_new_block and not self.sent_content_block_finish: - # Queue the sequence: content_block_stop -> content_block_start -> current_chunk + # Queue the sequence: content_block_stop -> content_block_start + # The trigger chunk itself is not emitted as a delta since the + # content_block_start already carries the relevant information. # 1. Stop current content block self.chunk_queue.append( @@ -284,9 +308,6 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): } ) - # 3. Queue the current chunk (don't lose it!) - self.chunk_queue.append(processed_chunk) - # Reset state for new block self.sent_content_block_finish = False diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index eeb55911ecf..2a2955fca37 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -43,8 +43,12 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): if "tool_choice" not in params: params.append("tool_choice") - # Only gpt-5.2 has been verified to support logprobs on Azure - if self.is_model_gpt_5_2_model(model): + # Only gpt-5.2 has been verified to support logprobs on Azure. + # The base OpenAI class includes logprobs for gpt-5.1+, but Azure + # hasn't verified support for gpt-5.1, so remove them unless gpt-5.2. + if self.is_model_gpt_5_1_model(model) and not self.is_model_gpt_5_2_model(model): + params = [p for p in params if p not in ["logprobs", "top_logprobs"]] + elif self.is_model_gpt_5_2_model(model): azure_supported_params = ["logprobs", "top_logprobs"] params.extend(azure_supported_params) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py index 4a170d666f5..eadc0da2f1f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py @@ -23,7 +23,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, ) -from litellm.types.utils import Delta, ModelResponse, StreamingChoices, Usage +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage class MockCompletionStreamWithContentAfterStopReason: @@ -32,16 +32,14 @@ class MockCompletionStreamWithContentAfterStopReason: def __init__(self): self.responses = [ # Initial text content - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content="Hello"), index=0, finish_reason=None ) ], ), - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content=" world"), index=0, finish_reason=None @@ -49,8 +47,7 @@ class MockCompletionStreamWithContentAfterStopReason: ], ), # Message delta with stop_reason AND usage (this is how it actually comes from the API) - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content=""), index=0, finish_reason="stop" @@ -60,8 +57,7 @@ class MockCompletionStreamWithContentAfterStopReason: ), # Additional content after the stop_reason - this simulates the scenario # where there might be additional content blocks after the main response - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content=" Additional content"), diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py index 9d4e58f3c88..1d25d719384 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py @@ -10,7 +10,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterato ) from litellm.types.utils import ( Delta, - ModelResponse, + ModelResponseStream, StreamingChoices, Usage, ChatCompletionDeltaToolCall, @@ -19,7 +19,7 @@ from litellm.types.utils import ( class MockCompletionStream: - def __init__(self, responses: List[ModelResponse]): + def __init__(self, responses: List[ModelResponseStream]): self.responses = responses self.index = 0 @@ -44,9 +44,8 @@ class MockCompletionStream: return response -def construct_text_chunk(text: str) -> ModelResponse: - return ModelResponse( - stream=True, +def construct_text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content=text), @@ -59,11 +58,10 @@ def construct_text_chunk(text: str) -> ModelResponse: def construct_split_tool_call( id: str, function_name: str, function_arg_parts: List[str] -) -> List[ModelResponse]: +) -> List[ModelResponseStream]: return [ # https://platform.openai.com/docs/guides/function-calling#streaming - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta( @@ -82,8 +80,7 @@ def construct_split_tool_call( ], ), *[ - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta( @@ -109,8 +106,7 @@ def construct_split_tool_call( def test_anthropic_stream_wrapper_single_tool_call(): responses = [ *construct_split_tool_call("tooluse_foo", "get_weather", ['{"city":', '"NY"}']), - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content="", stop_reason="tool_calls"), @@ -172,8 +168,7 @@ def test_anthropic_stream_wrapper_back_to_back_tool_calls(): responses = [ *construct_split_tool_call("tooluse_foo", "get_weather", ['{"city":', '"NY"}']), *construct_split_tool_call("tooluse_bar", "get_weather", ['{"city":', '"SF"}']), - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content="", stop_reason="tool_calls"), @@ -244,8 +239,7 @@ def test_anthropic_stream_wrapper_interleaved_tool_calls_and_text(): "tooluse_bar", "get_weather", ['{"city":', '"CHI"}'] ), construct_text_chunk("The weather is not so nice today."), - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content="", stop_reason="tool_calls"), diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py index dfcb9b3eb74..63fed907c3c 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py @@ -9,31 +9,28 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, ) -from litellm.types.utils import Delta, ModelResponse, StreamingChoices +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices # Create a simple test class MockCompletionStream: def __init__(self): self.responses = [ - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content="Hello"), index=0, finish_reason=None ) ], ), - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content=" World"), index=0, finish_reason=None ) ], ), - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content=""), index=0, finish_reason="stop" @@ -109,16 +106,14 @@ async def test_async_anthropic_sse_wrapper(): class AsyncMockCompletionStream: def __init__(self): self.responses = [ - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content="Hello"), index=0, finish_reason=None ) ], ), - ModelResponse( - stream=True, + ModelResponseStream( choices=[ StreamingChoices( delta=Delta(content=" World"), index=0, finish_reason=None From ab718444c57db2bc88f078d82a8821ca7f9ba7d7 Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Tue, 3 Mar 2026 18:28:34 -0300 Subject: [PATCH 098/147] Remove dead pending_new_content_block attribute Cleanup per review: this class attribute is no longer used after the __next__ refactor to queue-based approach. Co-Authored-By: Claude Opus 4.6 --- .../experimental_pass_through/adapters/streaming_iterator.py | 1 - 1 file changed, 1 deletion(-) 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 cdf8ac5ca82..7f17526e75c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -41,7 +41,6 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): type="text", text="", ) - pending_new_content_block: bool = False chunk_queue: deque = deque() # Queue for buffering multiple chunks def __init__( From d6ad312a4c2a2d14625210a9bd6daceb14940f97 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Mar 2026 13:37:11 -0800 Subject: [PATCH 099/147] fix(proxy): validate enforced_file_expires_after keys before access Add key validation for enforced_file_expires_after to return a clear 400 error instead of an unhandled KeyError 500. --- .../proxy/openai_files_endpoints/files_endpoints.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index a641fb1c69f..386ab2bf044 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -458,11 +458,22 @@ async def create_file( # noqa: PLR0915 team_metadata = user_api_key_dict.team_metadata or {} enforced_file_expiry = team_metadata.get("enforced_file_expires_after") if enforced_file_expiry is not None: + if "anchor" not in enforced_file_expiry or "seconds" not in enforced_file_expiry: + raise HTTPException( + status_code=400, + detail={ + "error": "enforced_file_expires_after must contain 'anchor' and 'seconds' keys", + }, + ) expires_after = FileExpiresAfter( anchor=enforced_file_expiry["anchor"], seconds=enforced_file_expiry["seconds"], ) + verbose_proxy_logger.info( + "create_file expires_after: %s", expires_after + ) + _create_file_request = CreateFileRequest( file=file_data, purpose=cast(CREATE_FILE_REQUESTS_PURPOSE, purpose), From 903ade4a1b3d10c83c2b7b91ffdc22322739e55e Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Mar 2026 13:51:41 -0800 Subject: [PATCH 100/147] fix(proxy): add anchor validation for file expiry, key validation for batch expiry Validate anchor is "created_at" in enforced_file_expires_after (matching user-provided path). Add key existence validation to batch endpoint for enforced_batch_output_expires_after. --- litellm/proxy/batches_endpoints/endpoints.py | 7 +++++++ litellm/proxy/openai_files_endpoints/files_endpoints.py | 9 ++++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 60905243369..850134b649c 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -125,6 +125,13 @@ async def create_batch( # noqa: PLR0915 "enforced_batch_output_expires_after" ) if enforced_batch_expiry is not None: + if "anchor" not in enforced_batch_expiry or "seconds" not in enforced_batch_expiry: + raise HTTPException( + status_code=400, + detail={ + "error": "enforced_batch_output_expires_after must contain 'anchor' and 'seconds' keys", + }, + ) _create_batch_data["output_expires_after"] = enforced_batch_expiry input_file_id = _create_batch_data.get("input_file_id", None) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 386ab2bf044..82cae8c64ea 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -465,8 +465,15 @@ async def create_file( # noqa: PLR0915 "error": "enforced_file_expires_after must contain 'anchor' and 'seconds' keys", }, ) + if enforced_file_expiry["anchor"] != "created_at": + raise HTTPException( + status_code=400, + detail={ + "error": f"enforced_file_expires_after anchor must be 'created_at', got '{enforced_file_expiry['anchor']}'", + }, + ) expires_after = FileExpiresAfter( - anchor=enforced_file_expiry["anchor"], + anchor="created_at", seconds=enforced_file_expiry["seconds"], ) From 657a60ea5b55fa3497e4222b12ec782fc54f2df8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 14:24:55 -0800 Subject: [PATCH 101/147] fix(audit): AND semantics for combined JSON filters; remove unused allTeams prop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix object_team_id + object_key_hash combining incorrectly as OR — each filter now adds an AND clause wrapping an internal OR over before_value and updated_values, so both conditions must be satisfied simultaneously - Rename helper to _build_json_field_or_condition to reflect its purpose - Remove allTeams from AuditLogsProps and its call site in index.tsx Co-Authored-By: Claude Sonnet 4.6 --- .../proxy/audit_logging_endpoints.py | 42 ++++++++++--------- .../src/components/view_logs/audit_logs.tsx | 2 - .../src/components/view_logs/index.tsx | 1 - 3 files changed, 22 insertions(+), 23 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index 5ab3669b50c..18ac29b9781 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -22,19 +22,25 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() -def _build_json_field_conditions( - field: str, json_key: str, value: str -) -> List[Dict[str, Any]]: +def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]: """ - Build OR conditions to match a value inside a JSON column at the given key. + Build an OR condition that matches a value inside a JSON column at the + given key, checking both before_value and updated_values. - Uses Prisma's JSON path filtering (PostgreSQL only). Returns a list of - two conditions — one for `before_value` and one for `updated_values` — to - be merged into the caller's top-level OR list. + Uses Prisma's JSON path filtering (PostgreSQL only). + + Example result (team_id="t1"): + {"OR": [ + {"before_value": {"path": ["team_id"], "string_contains": "t1"}}, + {"updated_values": {"path": ["team_id"], "string_contains": "t1"}}, + ]} """ - return [ - {field: {"path": [json_key], "string_contains": value}}, - ] + return { + "OR": [ + {"before_value": {"path": [json_key], "string_contains": value}}, + {"updated_values": {"path": [json_key], "string_contains": value}}, + ] + } @router.get( @@ -115,19 +121,15 @@ async def get_audit_logs( date_filter["lte"] = end_date where_conditions["updated_at"] = date_filter - # JSON field filters (PostgreSQL only) — search inside before_value and - # updated_values for a matching key/value pair. + # JSON field filters (PostgreSQL only) — each filter is AND'd with the + # others, but checks both before_value and updated_values internally (OR). if object_team_id: - where_conditions["OR"] = [ - *_build_json_field_conditions("before_value", "team_id", object_team_id), - *_build_json_field_conditions("updated_values", "team_id", object_team_id), + where_conditions["AND"] = where_conditions.get("AND", []) + [ + _build_json_field_or_condition("team_id", object_team_id) ] if object_key_hash: - existing_or: List[Dict[str, Any]] = where_conditions.get("OR", []) - where_conditions["OR"] = [ - *existing_or, - *_build_json_field_conditions("before_value", "token", object_key_hash), - *_build_json_field_conditions("updated_values", "token", object_key_hash), + where_conditions["AND"] = where_conditions.get("AND", []) + [ + _build_json_field_or_condition("token", object_key_hash) ] # Build sort conditions diff --git a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx index d3372eca6ff..693318acca8 100644 --- a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx @@ -7,7 +7,6 @@ import moment from "moment"; import { uiAuditLogsCall } from "../networking"; import { AuditLogEntry } from "./columns"; import { AuditLogDrawer } from "./AuditLogDrawer/AuditLogDrawer"; -import { Team } from "../key_team_helpers/key_list"; const { Search } = Input; @@ -18,7 +17,6 @@ interface AuditLogsProps { userID: string | null; isActive: boolean; premiumUser: boolean; - allTeams: Team[]; } const asset_logos_folder = "../ui/assets/"; diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index 22cc06a73e9..64dc53c0e77 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -714,7 +714,6 @@ export default function SpendLogsTable({ accessToken={accessToken} isActive={activeTab === "audit logs"} premiumUser={premiumUser} - allTeams={allTeams} /> From 35d2bc382f0185850fbe6c9e505a69d940fb1619 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Mar 2026 14:28:25 -0800 Subject: [PATCH 102/147] fix(batches): suppress PLR0915 lint for create_batch dispatch function --- litellm/batches/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index e73f73d2f33..e69c5a5c377 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -154,7 +154,7 @@ async def acreate_batch( @client -def create_batch( +def create_batch( # noqa: PLR0915 completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, From c8e6428eb77196030d65f68de9e8a09aab6ce72c Mon Sep 17 00:00:00 2001 From: ryan-crabbe <128659760+ryan-crabbe@users.noreply.github.com> Date: Tue, 3 Mar 2026 14:35:58 -0800 Subject: [PATCH 103/147] Update litellm/proxy/openai_files_endpoints/files_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/openai_files_endpoints/files_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 82cae8c64ea..44bd9b09d8e 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -477,7 +477,7 @@ async def create_file( # noqa: PLR0915 seconds=enforced_file_expiry["seconds"], ) - verbose_proxy_logger.info( + verbose_proxy_logger.debug( "create_file expires_after: %s", expires_after ) From 5da7fa9ac18a0b82a1a294389d6cf146fa443dd9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 14:39:39 -0800 Subject: [PATCH 104/147] [Feature] UI - Virtual Keys: Add manual spend reset to unblock keys Adds a "Reset Spend" button to the key detail view so proxy admins and team admins can immediately reset a key's spend to $0, unblocking keys that have hit their budget limit without waiting for the next scheduled budget reset. Co-Authored-By: Claude Sonnet 4.6 --- .../hooks/keys/useResetKeySpend.ts | 66 ++++++++ .../components/templates/KeyInfoHeader.tsx | 8 + .../templates/key_info_view.test.tsx | 141 ++++++++++++++++++ .../components/templates/key_info_view.tsx | 50 ++++++- 4 files changed, 264 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts new file mode 100644 index 00000000000..a845fc5881a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useResetKeySpend.ts @@ -0,0 +1,66 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { + getProxyBaseUrl, + getGlobalLitellmHeaderName, + deriveErrorMessage, + handleError, +} from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { keyKeys } from "./useKeys"; + +// ── Types ───────────────────────────────────────────────────────────────────── + +export interface ResetKeySpendResponse { + key_hash: string; + spend: number; + previous_spend: number; + max_budget: number | null; + budget_reset_at: string | null; +} + +// ── Fetch function ──────────────────────────────────────────────────────────── + +export const resetKeySpend = async ( + accessToken: string, + keyToken: string, +): Promise => { + const baseUrl = getProxyBaseUrl(); + const url = `${baseUrl ? `${baseUrl}/key/${keyToken}/reset_spend` : `/key/${keyToken}/reset_spend`}`; + + const response = await fetch(url, { + method: "POST", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ reset_to: 0 }), + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + return response.json(); +}; + +// ── Hook ────────────────────────────────────────────────────────────────────── + +export const useResetKeySpend = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (keyToken) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return resetKeySpend(accessToken, keyToken); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: keyKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx index 1befd657843..df3678afab6 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx @@ -11,6 +11,7 @@ import { ClockCircleOutlined, ThunderboltOutlined, SafetyCertificateOutlined, + DollarOutlined, } from "@ant-design/icons"; import LabeledField from "../common_components/LabeledField"; @@ -33,6 +34,7 @@ interface KeyInfoHeaderProps { onCreateNew?: () => void; onRegenerate?: () => void; onDelete?: () => void; + onResetSpend?: () => void; canModifyKey?: boolean; backButtonText?: string; regenerateDisabled?: boolean; @@ -45,6 +47,7 @@ export function KeyInfoHeader({ onCreateNew, onRegenerate, onDelete, + onResetSpend, canModifyKey = true, backButtonText = "Back to Keys", regenerateDisabled = false, @@ -77,6 +80,11 @@ export function KeyInfoHeader({
{canModifyKey && ( + {onResetSpend && ( + + )} )} diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 35a75b2f2fe..b77b13b1591 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -450,8 +450,8 @@ export default function KeyInfoView({ $0?

- Current spend: ${formatNumberWithCommas(currentKeyData.spend, 4)}. The key will be - immediately unblocked and able to make requests again. + Current spend: ${formatNumberWithCommas(currentKeyData.spend, 4)}. Spend history is + preserved in logs. This resets the current period spend counter, the same as an automatic budget reset.

From 5b2110ddb56492c6b8876c1e5cf3608d754d3820 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 15:16:47 -0800 Subject: [PATCH 112/147] Polish Reset Spend button and modal - Move Reset Spend button after Regenerate Key in header - Make modal OK button danger style with text "Reset" Co-Authored-By: Claude Sonnet 4.6 --- .../src/components/templates/KeyInfoHeader.tsx | 10 +++++----- .../src/components/templates/key_info_view.test.tsx | 4 ++-- .../src/components/templates/key_info_view.tsx | 3 ++- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx index c750b1b4ac7..93ebae9c4be 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx @@ -80,11 +80,6 @@ export function KeyInfoHeader({ {canModifyKey && ( - {onResetSpend && ( - - )} + {onResetSpend && ( + + )} diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index c19866e5d84..f269ad96a27 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -642,7 +642,7 @@ describe("KeyInfoView", () => { await waitFor(() => { expect(screen.getByText("Reset Key Spend")).toBeInTheDocument(); - expect(screen.getByText(/reset to \$0/i)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /^reset$/i })).toBeInTheDocument(); }); }); @@ -670,7 +670,7 @@ describe("KeyInfoView", () => { }); // Click the confirm button in the modal - await userEvent.click(screen.getByRole("button", { name: /reset to \$0/i })); + await userEvent.click(screen.getByRole("button", { name: /^reset$/i })); await waitFor(() => { expect(mockResetKeySpendMutate).toHaveBeenCalledWith( diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index b77b13b1591..6733cd6a595 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -442,7 +442,8 @@ export default function KeyInfoView({ open={isResetSpendModalOpen} onOk={handleResetSpend} onCancel={() => setIsResetSpendModalOpen(false)} - okText="Reset to $0" + okText="Reset" + okButtonProps={{ danger: true }} confirmLoading={resetSpendLoading} >

From 1a23ba41fd29823f8f845ddd6de83bf68700f498 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Tue, 3 Mar 2026 15:18:52 -0800 Subject: [PATCH 113/147] Fixed: undefined is not a valid Select option value --- ui/litellm-dashboard/src/components/model_info_view.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 0ba86e9f243..b8a5cb10fa1 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -627,7 +627,7 @@ export default function ModelInfoView({ : [], tags: Array.isArray(localModelData.litellm_params?.tags) ? localModelData.litellm_params.tags : [], health_check_model: isWildcardModel ? localModelData.model_info?.health_check_model : null, - litellm_credential_name: localModelData.litellm_params?.litellm_credential_name || undefined, + litellm_credential_name: localModelData.litellm_params?.litellm_credential_name || "", litellm_extra_params: JSON.stringify( Object.fromEntries( Object.entries(localModelData.litellm_params || {}).filter( @@ -970,7 +970,7 @@ export default function ModelInfoView({ (option?.label ?? "").toLowerCase().includes(input.toLowerCase()) } options={[ - { value: undefined, label: "None" }, + { value: "", label: "None" }, ...credentialsList.map((credential) => ({ value: credential.credential_name, label: credential.credential_name, From 98b9bc8b722b3f8d590d4cbe678280f8fa95cca7 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Wed, 4 Mar 2026 01:43:02 +0200 Subject: [PATCH 114/147] fix: resolve base_model in /cost/estimate for Azure custom deployments (#22724) The _resolve_model_for_cost_lookup function was only checking litellm_params.model when resolving model names from the router. For Azure custom deployment names (e.g. azure/openai/gpt-5.3-codex), this deployment name doesn't exist in the model cost map, so cost returned /bin/zsh. Now checks model_info.base_model and litellm_params.base_model first, falling back to litellm_params.model only if no base_model is set. This matches how the router resolves base_model everywhere else. --- .../cost_tracking_settings.py | 15 ++- .../test_cost_tracking_settings.py | 120 ++++++++++++++++++ 2 files changed, 133 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 6cdadfe216a..38dd4578c05 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -54,16 +54,27 @@ def _resolve_model_for_cost_lookup(model: str) -> Tuple[str, Optional[str]]: deployments = llm_router.get_model_list(model_name=model) if deployments and len(deployments) > 0: - # Get the first deployment's litellm model first_deployment = deployments[0] litellm_params = first_deployment.get("litellm_params", {}) + model_info = first_deployment.get("model_info", {}) + + # Check base_model first (needed for Azure custom deployment names) + base_model = model_info.get("base_model") or litellm_params.get( + "base_model" + ) + if base_model: + verbose_proxy_logger.debug( + f"Resolved model '{model}' to base_model '{base_model}' from router" + ) + custom_llm_provider = litellm_params.get("custom_llm_provider") + return base_model, custom_llm_provider + resolved_model = litellm_params.get("model") if resolved_model: verbose_proxy_logger.debug( f"Resolved model '{model}' to '{resolved_model}' from router" ) - # Extract custom_llm_provider if present custom_llm_provider = litellm_params.get("custom_llm_provider") return resolved_model, custom_llm_provider except Exception as e: diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index 275240dcc9e..1284cceba26 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -270,3 +270,123 @@ class TestCostTrackingSettings: assert "error" in response_data["detail"] assert "STORE_MODEL_IN_DB" in response_data["detail"]["error"] + + +class TestResolveModelForCostLookup: + """Tests for _resolve_model_for_cost_lookup base_model resolution.""" + + def test_resolves_base_model_for_azure_deployment(self): + """ + When a model group has base_model set in model_info, + _resolve_model_for_cost_lookup should return the base_model + instead of the raw litellm_params.model (Azure deployment name). + """ + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "gpt-5.3-codex", + "litellm_params": { + "model": "azure/openai/gpt-5.3-codex", + "api_base": "https://fake.openai.azure.com/", + "api_key": "fake-key", + }, + "model_info": { + "id": "test-id", + "base_model": "azure/gpt-4o", + }, + } + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + mock_router, + ): + resolved_model, provider = _resolve_model_for_cost_lookup("gpt-5.3-codex") + + assert resolved_model == "azure/gpt-4o" + mock_router.get_model_list.assert_called_once_with(model_name="gpt-5.3-codex") + + def test_falls_back_to_litellm_params_model_when_no_base_model(self): + """ + When no base_model is set, should fall back to litellm_params.model. + """ + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "openai/gpt-4", + }, + "model_info": { + "id": "test-id", + }, + } + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + mock_router, + ): + resolved_model, provider = _resolve_model_for_cost_lookup("gpt-4") + + assert resolved_model == "openai/gpt-4" + + def test_resolves_base_model_from_litellm_params(self): + """ + When base_model is in litellm_params (not model_info), + it should still be resolved. + """ + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "my-azure-model", + "litellm_params": { + "model": "azure/my-custom-deployment", + "base_model": "azure/gpt-4o-mini", + }, + "model_info": { + "id": "test-id", + }, + } + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + mock_router, + ): + resolved_model, provider = _resolve_model_for_cost_lookup( + "my-azure-model" + ) + + assert resolved_model == "azure/gpt-4o-mini" + + def test_returns_original_model_when_no_router(self): + """ + When no router is available, should return the original model name. + """ + from litellm.proxy.management_endpoints.cost_tracking_settings import ( + _resolve_model_for_cost_lookup, + ) + + with patch( + "litellm.proxy.proxy_server.llm_router", + None, + ): + resolved_model, provider = _resolve_model_for_cost_lookup( + "azure/openai/gpt-5.3-codex" + ) + + assert resolved_model == "azure/openai/gpt-5.3-codex" + assert provider is None From 0a1b2635d7f8c2e17f9c3c5732bb9b8f3dc23c08 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 15:54:34 -0800 Subject: [PATCH 115/147] fix: allow team admins to access /key/{key}/reset_spend route The route-level auth check was blocking internal_user role (team admins) from reaching /key/{key}/reset_spend because KEY_RESET_SPEND was missing from key_management_routes. Added it so team admins pass the route check and the endpoint's existing _check_proxy_or_team_admin_for_key enforces actual authorization. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/_types.py | 1 + .../proxy/auth/test_route_checks.py | 74 +++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5122c64ea64..8d49020461d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -512,6 +512,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.KEY_UNBLOCK.value, KeyManagementRoutes.KEY_BULK_UPDATE.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value, + KeyManagementRoutes.KEY_RESET_SPEND.value, ] management_routes = [ diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index ec1a13b8abc..f1e96f3e660 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -1116,3 +1116,77 @@ def test_route_in_additional_public_routes_exact_match(): assert route_in_additonal_public_routes("/status") is True # Non-matching routes should fail assert route_in_additonal_public_routes("/other") is False + + +def test_internal_user_can_access_key_reset_spend_route(): + """ + Regression test: team admins (role=internal_user) should pass the route-level + check for /key/{hash}/reset_spend. The endpoint itself enforces team admin status. + """ + user_obj = LiteLLM_UserTable( + user_id="team-admin-user", + user_email="teamadmin@example.com", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + valid_token = UserAPIKeyAuth( + user_id="team-admin-user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + key_hash = "baec26d2901589fe9fec76610e6e2be4895cdd8e19b3ada9a4fa2eb85e1901ae" + route = f"/key/{key_hash}/reset_spend" + + # Should not raise — the route-level check must pass for team admins + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + + +def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_reset_spend(): + """ + An internal_user passes the route check for /key/{hash}/reset_spend + (authorization is deferred to the endpoint), but is still blocked from + admin-only routes like /config/update. + """ + user_obj = LiteLLM_UserTable( + user_id="regular-user", + user_email="user@example.com", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + valid_token = UserAPIKeyAuth( + user_id="regular-user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + key_hash = "baec26d2901589fe9fec76610e6e2be4895cdd8e19b3ada9a4fa2eb85e1901ae" + + # /key/{hash}/reset_spend passes the route check for internal_user + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=f"/key/{key_hash}/reset_spend", + request=request, + valid_token=valid_token, + request_data={}, + ) + + # /config/update is still blocked + with pytest.raises(Exception) as exc_info: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/config/update", + request=request, + valid_token=valid_token, + request_data={}, + ) + assert "Only proxy admin can be used to generate" in str(exc_info.value) From 326ff422f1faad0a7654289d97824270be6c5686 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Mar 2026 15:54:46 -0800 Subject: [PATCH 116/147] fix(ui): Audit logs table and drawer polish Table: - Use DefaultProxyAdminTag for changed_by column - Remove tooltips on Object ID and API Key columns - Rename API Key column to API Key (Hash) - Move pagination controls to upper-right of filter bar; add icon-only refresh button Drawer: - Object ID is now copyable - API Key (Hash) is copyable and no longer truncated - Changed By uses DefaultProxyAdminTag - Expand JSON view boxes from max-h-72 to max-h-96 - Remove unnecessary vertical scrollbar (drop overflow-auto h-full from body div; use flex column layout so header and content flow naturally) Co-Authored-By: Claude Sonnet 4.6 --- .../AuditLogDrawer/AuditLogDrawer.tsx | 40 +++--- .../src/components/view_logs/audit_logs.tsx | 125 ++++++------------ 2 files changed, 62 insertions(+), 103 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx index 1e8dca34f22..de76709f5b2 100644 --- a/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx @@ -1,7 +1,10 @@ -import { Drawer, Tag, Tooltip } from "antd"; +import { Drawer, Tag, Typography } from "antd"; import { CloseOutlined } from "@ant-design/icons"; import moment from "moment"; import { AuditLogEntry } from "../columns"; +import DefaultProxyAdminTag from "../../common_components/DefaultProxyAdminTag"; + +const { Text } = Typography; interface AuditLogDrawerProps { open: boolean; @@ -26,7 +29,7 @@ const ACTION_COLOR: Record = { function JsonBlock({ value }: { value: Record }) { return ( - 

+    
       {JSON.stringify(value, null, 2)}
     

   );
@@ -92,7 +95,7 @@ function DiffSection({ log }: { log: AuditLogEntry }) {
         : { note: "No differing fields detected" };
   }
 
-  const renderValue = (value: Record | null | undefined, label: string) => {
+  const renderValue = (value: Record | null | undefined) => {
     if (!value || Object.keys(value).length === 0) {
       return 
N/A
;
     }
@@ -125,11 +128,11 @@ function DiffSection({ log }: { log: AuditLogEntry }) {
     

       

         
Before

-        {renderValue(displayBefore, "before")}
+        {renderValue(displayBefore)}
       

       

         
After

-        {renderValue(displayAfter, "after")}
+        {renderValue(displayAfter)}
       

     

   );
@@ -150,10 +153,10 @@ export function AuditLogDrawer({ open, onClose, log }: AuditLogDrawerProps) {
       closable={false}
       mask={true}
       maskClosable={true}
-      styles={{ body: { padding: 0 }, header: { display: "none" } }}
+      styles={{ body: { padding: 0, display: "flex", flexDirection: "column" }, header: { display: "none" } }}
     >
       {/* Header */}
-      

+      

         

           
             {log.action}
@@ -172,7 +175,7 @@ export function AuditLogDrawer({ open, onClose, log }: AuditLogDrawerProps) {
       

 
       {/* Body */}
-      

+      

         {/* Metadata */}
         

           

@@ -182,21 +185,22 @@ export function AuditLogDrawer({ open, onClose, log }: AuditLogDrawerProps) {
           
-                {log.object_id}
-              
+              
+                {log.object_id}
+              
             }
           />
-          
           }
+          />
+          
-                  
-                    {log.changed_by_api_key.slice(0, 12)}…
-                  
-                
+                
+                  {log.changed_by_api_key}
+                
               ) : (
                 "—"
               )
diff --git a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx
index 693318acca8..71300fe0427 100644
--- a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx
@@ -1,12 +1,13 @@
 import { useState } from "react";
 import { useQuery, keepPreviousData } from "@tanstack/react-query";
-import { Table, Tag, Input, Select, Button, Tooltip } from "antd";
+import { Table, Tag, Input, Select, Button, Pagination } from "antd";
 import { ReloadOutlined } from "@ant-design/icons";
-import type { ColumnsType, TablePaginationConfig } from "antd/es/table";
+import type { ColumnsType } from "antd/es/table";
 import moment from "moment";
 import { uiAuditLogsCall } from "../networking";
 import { AuditLogEntry } from "./columns";
 import { AuditLogDrawer } from "./AuditLogDrawer/AuditLogDrawer";
+import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag";
 
 const { Search } = Input;
 
@@ -97,14 +98,7 @@ export default function AuditLogs({
     placeholderData: keepPreviousData,
   });
 
-  const handleFilterChange = () => {
-    // Reset to page 1 whenever a filter changes
-    setPage(1);
-  };
-
-  const handleTableChange = (pagination: TablePaginationConfig) => {
-    setPage(pagination.current ?? 1);
-  };
+  const resetPage = () => setPage(1);
 
   const handleRowClick = (log: AuditLogEntry) => {
     setSelectedLog(log);
@@ -146,9 +140,7 @@ export default function AuditLogs({
       dataIndex: "object_id",
       key: "object_id",
       render: (val: string) => (
-        
-          {val}
-        
+        {val}
       ),
     },
     {
@@ -156,18 +148,16 @@ export default function AuditLogs({
       dataIndex: "changed_by",
       key: "changed_by",
       width: 200,
-      render: (val: string) => val || "—",
+      render: (val: string) => ,
     },
     {
-      title: "API Key",
+      title: "API Key (Hash)",
       dataIndex: "changed_by_api_key",
       key: "changed_by_api_key",
       width: 140,
       render: (val: string) =>
         val ? (
-          
-            {val.slice(0, 12)}…
-          
+          {val.slice(0, 12)}…
         ) : (
           "—"
         ),
@@ -212,76 +202,37 @@ export default function AuditLogs({
         

           

             
Audit Logs

-            
           

 
-          {/* Filters */}
-          

+          {/* Filters + pagination on same row */}
+          

              {
-                setObjectId(val);
-                handleFilterChange();
-              }}
-              onChange={(e) => {
-                if (!e.target.value) {
-                  setObjectId("");
-                  handleFilterChange();
-                }
-              }}
+              onSearch={(val) => { setObjectId(val); resetPage(); }}
+              onChange={(e) => { if (!e.target.value) { setObjectId(""); resetPage(); } }}
             />
              {
-                setChangedBy(val);
-                handleFilterChange();
-              }}
-              onChange={(e) => {
-                if (!e.target.value) {
-                  setChangedBy("");
-                  handleFilterChange();
-                }
-              }}
+              onSearch={(val) => { setChangedBy(val); resetPage(); }}
+              onChange={(e) => { if (!e.target.value) { setChangedBy(""); resetPage(); } }}
             />
              {
-                setTeamId(val);
-                handleFilterChange();
-              }}
-              onChange={(e) => {
-                if (!e.target.value) {
-                  setTeamId("");
-                  handleFilterChange();
-                }
-              }}
+              onSearch={(val) => { setTeamId(val); resetPage(); }}
+              onChange={(e) => { if (!e.target.value) { setTeamId(""); resetPage(); } }}
             />
              {
-                setKeyHash(val);
-                handleFilterChange();
-              }}
-              onChange={(e) => {
-                if (!e.target.value) {
-                  setKeyHash("");
-                  handleFilterChange();
-                }
-              }}
+              onSearch={(val) => { setKeyHash(val); resetPage(); }}
+              onChange={(e) => { if (!e.target.value) { setKeyHash(""); resetPage(); } }}
             />
              {
-                setTableName(val);
-                handleFilterChange();
-              }}
+              onChange={(val) => { setTableName(val); resetPage(); }}
             />
+
+            {/* Pagination + refresh pushed to the right */}
+            

+              

           

         

 
-        {/* Table */}
+        {/* Table — pagination handled in header */}
         
           columns={columns}
           dataSource={auditLogs}
           rowKey="id"
           loading={query.isLoading}
           size="small"
+          pagination={false}
           onRow={(record) => ({
             onClick: () => handleRowClick(record),
             style: { cursor: "pointer" },
           })}
-          pagination={{
-            current: page,
-            pageSize: PAGE_SIZE,
-            total,
-            showTotal: (t) => `${t} total`,
-            showSizeChanger: false,
-            onChange: (p) => setPage(p),
-          }}
-          onChange={handleTableChange}
         />
       

 

From 51f18b05a97f9352ef6fa20a064fb253c03565f1 Mon Sep 17 00:00:00 2001
From: Alejandro Tapia <67175024+atapia27@users.noreply.github.com>
Date: Tue, 3 Mar 2026 16:00:16 -0800
Subject: [PATCH 117/147] Update
 ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 .../src/components/team/TeamVirtualKeysTable.tsx                | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
index de3061a003d..c3cdd7fba5b 100644
--- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
@@ -569,7 +569,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
     onSortingChange: handleSortingChange,
     onPaginationChange: setTablePagination,
     getCoreRowModel: getCoreRowModel(),
-    getSortedRowModel: getSortedRowModel(),
+    // getSortedRowModel not needed — manualSorting: true delegates sorting to the server
     enableSorting: true,
     manualSorting: true, // Server sorts via useKeys. Avoid redundant client-side sort
     manualPagination: true,

From 5fa0c6e994de825502c3d5b1f8278a5ef03a4850 Mon Sep 17 00:00:00 2001
From: yuneng-jiang 
Date: Tue, 3 Mar 2026 16:00:46 -0800
Subject: [PATCH 118/147] fix(ui): Copyable JSON blocks in audit drawer; custom
 table spinner
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Replace plain JsonBlock with CopyableJsonBlock: header row with label
  and copy icon (CheckOutlined on success), pre block below — matches
  the spend logs RequestResponsePanel pattern
- Key-table rows use the same card chrome for visual consistency
- Remove separate "Changes" section label (now redundant with block headers)
- Table loading spinner replaced with custom Spin + LoadingOutlined

Co-Authored-By: Claude Sonnet 4.6 
---
 .../AuditLogDrawer/AuditLogDrawer.tsx         | 105 ++++++++++++------
 .../src/components/view_logs/audit_logs.tsx   |   9 +-
 2 files changed, 78 insertions(+), 36 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx
index de76709f5b2..19989ef4882 100644
--- a/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx
@@ -1,5 +1,6 @@
 import { Drawer, Tag, Typography } from "antd";
-import { CloseOutlined } from "@ant-design/icons";
+import { CloseOutlined, CopyOutlined, CheckOutlined } from "@ant-design/icons";
+import { useState, useCallback } from "react";
 import moment from "moment";
 import { AuditLogEntry } from "../columns";
 import DefaultProxyAdminTag from "../../common_components/DefaultProxyAdminTag";
@@ -27,11 +28,48 @@ const ACTION_COLOR: Record = {
   rotated: "orange",
 };
 
-function JsonBlock({ value }: { value: Record }) {
+function CopyableJsonBlock({ label, value }: { label: string; value: Record }) {
+  const [copied, setCopied] = useState(false);
+
+  const handleCopy = useCallback(async () => {
+    try {
+      const text = JSON.stringify(value, null, 2);
+      if (navigator.clipboard && window.isSecureContext) {
+        await navigator.clipboard.writeText(text);
+      } else {
+        const el = document.createElement("textarea");
+        el.value = text;
+        el.style.position = "fixed";
+        el.style.opacity = "0";
+        document.body.appendChild(el);
+        el.focus();
+        el.select();
+        document.execCommand("copy");
+        document.body.removeChild(el);
+      }
+      setCopied(true);
+      setTimeout(() => setCopied(false), 2000);
+    } catch (e) {
+      console.error("Copy failed:", e);
+    }
+  }, [value]);
+
   return (
-    
-      {JSON.stringify(value, null, 2)}
-    

+    

+      

+        {label}
+        
+      

+      
+        {JSON.stringify(value, null, 2)}
+      

+    

   );
 }
 
@@ -95,45 +133,51 @@ function DiffSection({ log }: { log: AuditLogEntry }) {
         : { note: "No differing fields detected" };
   }
 
-  const renderValue = (value: Record | null | undefined) => {
+  const renderValue = (label: string, value: Record | null | undefined) => {
     if (!value || Object.keys(value).length === 0) {
-      return 
N/A
;
+      return (
+        

+          

+            {label}
+          

+          
N/A

+        

+      );
     }
 
-    // For key table updates, filter to only show meaningful fields
+    // For key table updates, show only meaningful fields as plain text
     if (isKeyTable && isUpdateAction) {
       const knownKeyFields = ["token", "spend", "max_budget"];
       const hasOnlyKnown = Object.keys(value).every((k) => knownKeyFields.includes(k));
       if (hasOnlyKnown && !("note" in value)) {
         return (
-          

-            {value.token !== undefined && (
-              
Token: {value.token ?? "N/A"}

-            )}
-            {value.spend !== undefined && (
-              
Spend: ${Number(value.spend).toFixed(6)}

-            )}
-            {value.max_budget !== undefined && (
-              
Max Budget: ${Number(value.max_budget).toFixed(6)}

-            )}
+          

+            

+              {label}
+            

+            

+              {value.token !== undefined && (
+                
Token: {value.token ?? "N/A"}

+              )}
+              {value.spend !== undefined && (
+                
Spend: ${Number(value.spend).toFixed(6)}

+              )}
+              {value.max_budget !== undefined && (
+                
Max Budget: ${Number(value.max_budget).toFixed(6)}

+              )}
+            

           

         );
       }
     }
 
-    return ;
+    return ;
   };
 
   return (
     

-      

-        
Before

-        {renderValue(displayBefore)}
-      

-      

-        
After

-        {renderValue(displayAfter)}
-      

+      {renderValue("Before", displayBefore)}
+      {renderValue("After", displayAfter)}
     

   );
 }
@@ -209,12 +253,7 @@ export function AuditLogDrawer({ open, onClose, log }: AuditLogDrawerProps) {
         

 
         {/* Diff */}
-        

-          

-            Changes
-          

-          
-        

+        
       

     
   );
diff --git a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx
index 71300fe0427..b16ba30049d 100644
--- a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx
@@ -1,7 +1,7 @@
 import { useState } from "react";
 import { useQuery, keepPreviousData } from "@tanstack/react-query";
-import { Table, Tag, Input, Select, Button, Pagination } from "antd";
-import { ReloadOutlined } from "@ant-design/icons";
+import { Table, Tag, Input, Select, Button, Pagination, Spin } from "antd";
+import { ReloadOutlined, LoadingOutlined } from "@ant-design/icons";
 import type { ColumnsType } from "antd/es/table";
 import moment from "moment";
 import { uiAuditLogsCall } from "../networking";
@@ -285,7 +285,10 @@ export default function AuditLogs({
           columns={columns}
           dataSource={auditLogs}
           rowKey="id"
-          loading={query.isLoading}
+          loading={{
+            spinning: query.isLoading,
+            indicator: } size="small" />,
+          }}
           size="small"
           pagination={false}
           onRow={(record) => ({

From 383f79f36f4a9a6cfe38bfe786a5a940f528325d Mon Sep 17 00:00:00 2001
From: Alejandro Tapia <67175024+atapia27@users.noreply.github.com>
Date: Tue, 3 Mar 2026 16:00:57 -0800
Subject: [PATCH 119/147] Update
 ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 .../src/components/key_team_helpers/filter_helpers.ts           | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
index 97166f92735..a1602fc3aea 100644
--- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
@@ -1,4 +1,4 @@
-import { teamListCall, organizationListCall, keyAliasesCall, keyListCall } from "../networking";
+import { teamListCall, organizationListCall, keyListCall } from "../networking";
 import { Team } from "./key_list";
 import { Organization } from "../networking";
 

From 47ff0d1d863426a3390dfb30bd871b98bf28b9a1 Mon Sep 17 00:00:00 2001
From: Alejandro Tapia 
Date: Tue, 3 Mar 2026 16:03:25 -0800
Subject: [PATCH 120/147] Fixed: 50 sequential API calls on mount is excessive

---
 .../key_team_helpers/filter_helpers.ts        | 109 +++++++++++-------
 1 file changed, 68 insertions(+), 41 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
index a1602fc3aea..539ac172569 100644
--- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
@@ -9,12 +9,35 @@ export interface TeamFilterOptions {
 }
 
 const FILTER_OPTIONS_PAGE_SIZE = 100; // API max per page
-const MAX_PAGES = 50; // Cap at 5000 keys to avoid unbounded fetches
+const MAX_PAGES = 10; // Cap at 1000 keys; filter completeness beyond ~500 has diminishing returns
+
+const processKeysIntoOptions = (
+  keys: Array>,
+  keyAliases: Set,
+  organizationIds: Set,
+  userMap: Map,
+) => {
+  for (const key of keys) {
+    const alias = key?.key_alias;
+    if (alias && typeof alias === "string") {
+      keyAliases.add(alias.trim());
+    }
+    const orgId = key?.organization_id;
+    if (orgId && typeof orgId === "string") {
+      organizationIds.add(orgId.trim());
+    }
+    const userId = key?.user_id;
+    if (userId && typeof userId === "string") {
+      const email = (key?.user as { user_email?: string })?.user_email || userId;
+      userMap.set(userId, email);
+    }
+  }
+};
 
 /**
- * Fetches filter options (key aliases, org IDs, user IDs) from all team keys.
- * Paginates through pages to build complete dropdowns. Capped at 50 pages
- * (5000 keys) to limit load for very large teams.
+ * Fetches filter options (key aliases, org IDs, user IDs) from team keys.
+ * Fetches page 1 first to get totalPages, then batches remaining pages with
+ * Promise.all. Capped at 10 pages (1000 keys)
  */
 export const fetchTeamFilterOptions = async (
   accessToken: string | null,
@@ -29,46 +52,50 @@ export const fetchTeamFilterOptions = async (
     const organizationIds = new Set();
     const userMap = new Map();
 
-    let page = 1;
-    let totalPages = 1;
+    // First request: get page 1 and totalPages
+    const firstResponse = await keyListCall(
+      accessToken,
+      null,
+      teamId,
+      null,
+      null,
+      null,
+      1,
+      FILTER_OPTIONS_PAGE_SIZE,
+      null,
+      null,
+      "user",
+      null,
+    );
 
-    do {
-      const response = await keyListCall(
-        accessToken,
-        null,
-        teamId,
-        null,
-        null,
-        null,
-        page,
-        FILTER_OPTIONS_PAGE_SIZE,
-        null,
-        null,
-        "user",
-        null,
+    const firstKeys = firstResponse?.keys || [];
+    const totalPages = firstResponse?.total_pages ?? 1;
+    processKeysIntoOptions(firstKeys, keyAliases, organizationIds, userMap);
+
+    // Batch fetch remaining pages (2 through min(totalPages, MAX_PAGES)) in parallel
+    const pagesToFetch = Math.min(totalPages, MAX_PAGES) - 1;
+    if (pagesToFetch > 0) {
+      const pagePromises = Array.from({ length: pagesToFetch }, (_, i) =>
+        keyListCall(
+          accessToken,
+          null,
+          teamId,
+          null,
+          null,
+          null,
+          i + 2,
+          FILTER_OPTIONS_PAGE_SIZE,
+          null,
+          null,
+          "user",
+          null,
+        ),
       );
-
-      const keys = response?.keys || [];
-      totalPages = response?.total_pages ?? 1;
-
-      for (const key of keys) {
-        const alias = key?.key_alias;
-        if (alias && typeof alias === "string") {
-          keyAliases.add(alias.trim());
-        }
-        const orgId = key?.organization_id;
-        if (orgId && typeof orgId === "string") {
-          organizationIds.add(orgId.trim());
-        }
-        const userId = key?.user_id;
-        if (userId && typeof userId === "string") {
-          const email = key?.user?.user_email || userId;
-          userMap.set(userId, email);
-        }
+      const responses = await Promise.all(pagePromises);
+      for (const response of responses) {
+        processKeysIntoOptions(response?.keys || [], keyAliases, organizationIds, userMap);
       }
-
-      page++;
-    } while (page <= totalPages && page <= MAX_PAGES);
+    }
 
     return {
       keyAliases: Array.from(keyAliases).sort(),

From 34acbb39da74f890f22f014d3dcac431705efa45 Mon Sep 17 00:00:00 2001
From: Alejandro Tapia <67175024+atapia27@users.noreply.github.com>
Date: Tue, 3 Mar 2026 16:19:32 -0800
Subject: [PATCH 121/147] Update
 ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 .../src/components/team/TeamVirtualKeysTable.tsx                | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
index c3cdd7fba5b..bb4ab2b974d 100644
--- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
@@ -8,7 +8,7 @@ import {
   ColumnDef,
   flexRender,
   getCoreRowModel,
-  getSortedRowModel,
+  getCoreRowModel,
   PaginationState,
   SortingState,
   useReactTable,

From 18a2213cdb375d31f4cb255533c156d2452b7900 Mon Sep 17 00:00:00 2001
From: Alejandro Tapia 
Date: Tue, 3 Mar 2026 16:21:02 -0800
Subject: [PATCH 122/147] Fix: Promise.all loses all results on single page
 failure

---
 .../src/components/key_team_helpers/filter_helpers.ts  | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
index 539ac172569..b587e090d33 100644
--- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts
@@ -37,7 +37,7 @@ const processKeysIntoOptions = (
 /**
  * Fetches filter options (key aliases, org IDs, user IDs) from team keys.
  * Fetches page 1 first to get totalPages, then batches remaining pages with
- * Promise.all. Capped at 10 pages (1000 keys)
+ * Promise.allSettled (preserves successful pages if some fail). Capped at 10 pages (1000 keys)
  */
 export const fetchTeamFilterOptions = async (
   accessToken: string | null,
@@ -91,9 +91,11 @@ export const fetchTeamFilterOptions = async (
           null,
         ),
       );
-      const responses = await Promise.all(pagePromises);
-      for (const response of responses) {
-        processKeysIntoOptions(response?.keys || [], keyAliases, organizationIds, userMap);
+      const results = await Promise.allSettled(pagePromises);
+      for (const result of results) {
+        if (result.status === "fulfilled") {
+          processKeysIntoOptions(result.value?.keys || [], keyAliases, organizationIds, userMap);
+        }
       }
     }
 

From 5f38c3bfecd403bd126e6999d1729471a2f72701 Mon Sep 17 00:00:00 2001
From: Alejandro Tapia <67175024+atapia27@users.noreply.github.com>
Date: Tue, 3 Mar 2026 16:30:10 -0800
Subject: [PATCH 123/147] Update
 ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 .../src/components/team/TeamVirtualKeysTable.tsx                 | 1 -
 1 file changed, 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
index bb4ab2b974d..5d76b99ef91 100644
--- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
@@ -8,7 +8,6 @@ import {
   ColumnDef,
   flexRender,
   getCoreRowModel,
-  getCoreRowModel,
   PaginationState,
   SortingState,
   useReactTable,

From 90eb6729d57c269b717f886479e4e53d4a8979dd Mon Sep 17 00:00:00 2001
From: Krish Dholakia 
Date: Tue, 3 Mar 2026 18:19:12 -0800
Subject: [PATCH 124/147] Agent Tracing - support context_id based trace id
 propogation + nested llm calls  (#22626)

* style(ui/): distinguish agent calls from llm calls on ui

* feat: initial grouping working

* feat: set stable contextid for a2a calls - allows for easily passing to downstream llm/mcp calls

* feat(a2a_endpoints.py): fix tracing to avoid recreating logging objects for the same call

allows stable trace id usage

* fix(guardrail_endpoints): handle string ui_type values in _build_field_dict

_build_field_dict unconditionally called .value on ui_type, which crashes
for guardrail configs that use plain strings (e.g. BlockCodeExecutionGuardrailConfigModel
uses "multiselect" and "percentage"). Now checks with hasattr before calling .value.

Co-Authored-By: Claude Opus 4.6 

* fix: propagate trace/session id from headers in MCP server calls

Cherry-picked mcp_server/server.py fixes from 6feb9bab: adds
get_chain_id_from_headers to extract x-litellm-trace-id /
x-litellm-session-id from raw headers, and uses it in call_tool
and list_tools to keep spend logs and tracing consistent with A2A.

Co-Authored-By: Claude Opus 4.6 

---------

Co-authored-by: Claude Opus 4.6 
---
 docs/my-website/docs/response_api.md          |   2 +-
 litellm/a2a_protocol/main.py                  |  91 ++++--
 .../litellm_core_utils/get_litellm_params.py  |   8 +-
 litellm/litellm_core_utils/litellm_logging.py | 306 +++++++++---------
 .../proxy/_experimental/mcp_server/server.py  |  17 +-
 .../out/{404.html => 404/index.html}          |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../{budgets.html => budgets/index.html}      |   0
 .../{caching.html => caching/index.html}      |   0
 .../index.html}                               |   0
 .../{old-usage.html => old-usage/index.html}  |   0
 .../{prompts.html => prompts/index.html}      |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../out/{login.html => login/index.html}      |   0
 .../out/{logs.html => logs/index.html}        |   0
 .../{callback.html => callback/index.html}    |   0
 .../{model-hub.html => model-hub/index.html}  |   0
 .../{model_hub.html => model_hub/index.html}  |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../{policies.html => policies/index.html}    |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../{ui-theme.html => ui-theme/index.html}    |   0
 .../out/{teams.html => teams/index.html}      |   0
 .../{test-key.html => test-key/index.html}    |   0
 .../index.html}                               |   0
 .../index.html}                               |   0
 .../out/{usage.html => usage/index.html}      |   0
 .../out/{users.html => users/index.html}      |   0
 .../index.html}                               |   0
 .../proxy/agent_endpoints/a2a_endpoints.py    |  49 ++-
 .../proxy/guardrails/guardrail_endpoints.py   |   8 +-
 litellm/proxy/litellm_pre_call_utils.py       |  79 +++--
 .../spend_tracking/spend_tracking_utils.py    |  83 +++--
 litellm/utils.py                              |  52 +--
 .../guardrails/test_guardrail_endpoints.py    |  43 ++-
 .../proxy/test_litellm_pre_call_utils.py      | 100 ++++--
 .../LogDetailsDrawer/LogDetailsDrawer.tsx     |  39 ++-
 .../src/components/view_logs/TypeBadges.tsx   |  21 +-
 .../src/components/view_logs/columns.tsx      |  25 +-
 .../src/components/view_logs/constants.ts     |   3 +
 .../src/components/view_logs/index.tsx        |   9 +-
 ui/litellm-dashboard/tsconfig.json            |   2 +-
 51 files changed, 579 insertions(+), 358 deletions(-)
 rename litellm/proxy/_experimental/out/{404.html => 404/index.html} (100%)
 rename litellm/proxy/_experimental/out/{_not-found.html => _not-found/index.html} (100%)
 rename litellm/proxy/_experimental/out/{api-reference.html => api-reference/index.html} (100%)
 rename litellm/proxy/_experimental/out/experimental/{api-playground.html => api-playground/index.html} (100%)
 rename litellm/proxy/_experimental/out/experimental/{budgets.html => budgets/index.html} (100%)
 rename litellm/proxy/_experimental/out/experimental/{caching.html => caching/index.html} (100%)
 rename litellm/proxy/_experimental/out/experimental/{claude-code-plugins.html => claude-code-plugins/index.html} (100%)
 rename litellm/proxy/_experimental/out/experimental/{old-usage.html => old-usage/index.html} (100%)
 rename litellm/proxy/_experimental/out/experimental/{prompts.html => prompts/index.html} (100%)
 rename litellm/proxy/_experimental/out/experimental/{tag-management.html => tag-management/index.html} (100%)
 rename litellm/proxy/_experimental/out/{guardrails.html => guardrails/index.html} (100%)
 rename litellm/proxy/_experimental/out/{login.html => login/index.html} (100%)
 rename litellm/proxy/_experimental/out/{logs.html => logs/index.html} (100%)
 rename litellm/proxy/_experimental/out/mcp/oauth/{callback.html => callback/index.html} (100%)
 rename litellm/proxy/_experimental/out/{model-hub.html => model-hub/index.html} (100%)
 rename litellm/proxy/_experimental/out/{model_hub.html => model_hub/index.html} (100%)
 rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%)
 rename litellm/proxy/_experimental/out/{models-and-endpoints.html => models-and-endpoints/index.html} (100%)
 rename litellm/proxy/_experimental/out/{onboarding.html => onboarding/index.html} (100%)
 rename litellm/proxy/_experimental/out/{organizations.html => organizations/index.html} (100%)
 rename litellm/proxy/_experimental/out/{playground.html => playground/index.html} (100%)
 rename litellm/proxy/_experimental/out/{policies.html => policies/index.html} (100%)
 rename litellm/proxy/_experimental/out/settings/{admin-settings.html => admin-settings/index.html} (100%)
 rename litellm/proxy/_experimental/out/settings/{logging-and-alerts.html => logging-and-alerts/index.html} (100%)
 rename litellm/proxy/_experimental/out/settings/{router-settings.html => router-settings/index.html} (100%)
 rename litellm/proxy/_experimental/out/settings/{ui-theme.html => ui-theme/index.html} (100%)
 rename litellm/proxy/_experimental/out/{teams.html => teams/index.html} (100%)
 rename litellm/proxy/_experimental/out/{test-key.html => test-key/index.html} (100%)
 rename litellm/proxy/_experimental/out/tools/{mcp-servers.html => mcp-servers/index.html} (100%)
 rename litellm/proxy/_experimental/out/tools/{vector-stores.html => vector-stores/index.html} (100%)
 rename litellm/proxy/_experimental/out/{usage.html => usage/index.html} (100%)
 rename litellm/proxy/_experimental/out/{users.html => users/index.html} (100%)
 rename litellm/proxy/_experimental/out/{virtual-keys.html => virtual-keys/index.html} (100%)

diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md
index a7cf61ef16a..76899a17ccb 100644
--- a/docs/my-website/docs/response_api.md
+++ b/docs/my-website/docs/response_api.md
@@ -930,7 +930,7 @@ For Responses API with load balancing across deployments with **different API ke
 
 Notes:
 - User-key affinity is keyed on `metadata.user_api_key_hash` (the API key hash). The OpenAI `user` request parameter is an end-user identifier and is intentionally not used for deployment affinity.
-- Session-ID affinity is keyed on `metadata.session_id`. For proxy requests, this can be passed via the `x-litellm-session-id` HTTP header. For Python SDK requests, you can pass it via `litellm_metadata={"session_id": "value"}` in request args.
+- Session-ID affinity is keyed on `metadata.session_id`. For proxy requests, this can be passed via the `x-litellm-session-id` or `x-litellm-trace-id` HTTP header (they are interchangeable for call chaining). For Python SDK requests, you can pass it via `litellm_metadata={"session_id": "value"}` in request args.
 - `user_api_key_hash` is already SHA-256, and is used as-is (no double hashing).
 - Affinity is scoped by a stable model identifier (the model-map key, e.g. `model_map_information.model_map_key`) so model aliases map to the same stickiness bucket.
 - The mapping TTL is controlled by `deployment_affinity_ttl_seconds` (configured on Router init / proxy startup).
diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py
index 642dfaf023c..401b602fef5 100644
--- a/litellm/a2a_protocol/main.py
+++ b/litellm/a2a_protocol/main.py
@@ -24,11 +24,7 @@ from litellm.utils import client
 
 if TYPE_CHECKING:
     from a2a.client import A2AClient as A2AClientType
-    from a2a.types import (
-        AgentCard,
-        SendMessageRequest,
-        SendStreamingMessageRequest,
-    )
+    from a2a.types import AgentCard, SendMessageRequest, SendStreamingMessageRequest
 
 # Runtime imports with availability check
 A2A_SDK_AVAILABLE = False
@@ -124,13 +120,48 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str:
         litellm_logging_obj.model = model
         litellm_logging_obj.custom_llm_provider = custom_llm_provider
         litellm_logging_obj.model_call_details["model"] = model
-        litellm_logging_obj.model_call_details[
-            "custom_llm_provider"
-        ] = custom_llm_provider
+        litellm_logging_obj.model_call_details["custom_llm_provider"] = (
+            custom_llm_provider
+        )
 
     return agent_name
 
 
+async def _send_message_via_completion_bridge(
+    request: "SendMessageRequest",
+    custom_llm_provider: str,
+    api_base: Optional[str],
+    litellm_params: Dict[str, Any],
+) -> LiteLLMSendMessageResponse:
+    """
+    Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
+
+    Requires request; api_base is optional for providers that derive endpoint from model.
+    """
+    verbose_logger.info(
+        f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}"
+    )
+
+    from litellm.a2a_protocol.litellm_completion_bridge.handler import (
+        A2ACompletionBridgeHandler,
+    )
+
+    params = (
+        request.params.model_dump(mode="json")
+        if hasattr(request.params, "model_dump")
+        else dict(request.params)
+    )
+
+    response_dict = await A2ACompletionBridgeHandler.handle_non_streaming(
+        request_id=str(request.id),
+        params=params,
+        litellm_params=litellm_params,
+        api_base=api_base,
+    )
+
+    return LiteLLMSendMessageResponse.from_dict(response_dict)
+
+
 @client
 async def asend_message(
     a2a_client: Optional["A2AClientType"] = None,
@@ -193,39 +224,21 @@ async def asend_message(
         ```
     """
     litellm_params = litellm_params or {}
+    logging_obj = kwargs.get("litellm_logging_obj")
+    trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None
     custom_llm_provider = litellm_params.get("custom_llm_provider")
 
     # Route through completion bridge if custom_llm_provider is set
     if custom_llm_provider:
         if request is None:
             raise ValueError("request is required for completion bridge")
-        # api_base is optional for providers that derive endpoint from model (e.g., bedrock/agentcore)
-
-        verbose_logger.info(
-            f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}"
-        )
-
-        from litellm.a2a_protocol.litellm_completion_bridge.handler import (
-            A2ACompletionBridgeHandler,
-        )
-
-        # Extract params from request
-        params = (
-            request.params.model_dump(mode="json")
-            if hasattr(request.params, "model_dump")
-            else dict(request.params)
-        )
-
-        response_dict = await A2ACompletionBridgeHandler.handle_non_streaming(
-            request_id=str(request.id),
-            params=params,
-            litellm_params=litellm_params,
+        return await _send_message_via_completion_bridge(
+            request=request,
+            custom_llm_provider=custom_llm_provider,
             api_base=api_base,
+            litellm_params=litellm_params,
         )
 
-        # Convert to LiteLLMSendMessageResponse
-        return LiteLLMSendMessageResponse.from_dict(response_dict)
-
     # Standard A2A client flow
     if request is None:
         raise ValueError("request is required")
@@ -236,11 +249,13 @@ async def asend_message(
             raise ValueError(
                 "Either a2a_client or api_base is required for standard A2A flow"
             )
-        trace_id = str(uuid.uuid4())
+        trace_id = trace_id or str(uuid.uuid4())
         extra_headers = {"X-LiteLLM-Trace-Id": trace_id}
         if agent_id:
             extra_headers["X-LiteLLM-Agent-Id"] = agent_id
-        a2a_client = await create_a2a_client(base_url=api_base, extra_headers=extra_headers)
+        a2a_client = await create_a2a_client(
+            base_url=api_base, extra_headers=extra_headers
+        )
 
     # Type assertion: a2a_client is guaranteed to be non-None here
     assert a2a_client is not None
@@ -255,6 +270,10 @@ async def asend_message(
     )
     card_url = getattr(agent_card, "url", None) if agent_card else None
 
+    context_id = trace_id or str(uuid.uuid4())
+    if request.params.message.context_id is None:
+        request.params.message.context_id = context_id
+
     # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
     a2a_response = None
     for _ in range(2):  # max 2 attempts: original + 1 retry
@@ -606,7 +625,9 @@ async def create_a2a_client(
 
     if extra_headers:
         httpx_client.headers.update(extra_headers)
-        verbose_proxy_logger.debug(f"A2A client created with extra_headers={extra_headers}")
+        verbose_proxy_logger.debug(
+            f"A2A client created with extra_headers={extra_headers}"
+        )
 
     # Resolve agent card
     resolver = A2ACardResolver(
diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py
index 36a8dfdb5a6..c91e4b6de1d 100644
--- a/litellm/litellm_core_utils/get_litellm_params.py
+++ b/litellm/litellm_core_utils/get_litellm_params.py
@@ -1,6 +1,5 @@
 from typing import Optional
 
-
 # Pre-define optional kwargs keys as frozenset for O(1) lookups
 # These are extracted from kwargs only if present, avoiding unnecessary .get() calls
 _OPTIONAL_KWARGS_KEYS = frozenset({
@@ -95,6 +94,13 @@ def get_litellm_params(
     litellm_request_debug: Optional[bool] = None,
     **kwargs,
 ) -> dict:
+    # Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining)
+    _meta = metadata or {}
+    if litellm_session_id is None:
+        litellm_session_id = _meta.get("session_id") or _meta.get("trace_id")
+    if litellm_trace_id is None:
+        litellm_trace_id = _meta.get("trace_id") or _meta.get("session_id")
+
     # Build base dict with explicit parameters (always included)
     litellm_params = {
         "acompletion": acompletion,
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 5e5a6cea1b2..6f587abcdf1 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -133,8 +133,8 @@ from ..integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
 from ..integrations.azure_storage.azure_storage import AzureBlobStorageLogger
 from ..integrations.custom_prompt_management import CustomPromptManagement
 from ..integrations.datadog.datadog import DataDogLogger
-from ..integrations.datadog.datadog_metrics import DatadogMetricsLogger
 from ..integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
+from ..integrations.datadog.datadog_metrics import DatadogMetricsLogger
 from ..integrations.dotprompt import DotpromptManager
 from ..integrations.dynamodb import DyanmoDBLogger
 from ..integrations.galileo import GalileoObserve
@@ -352,9 +352,9 @@ class Logging(LiteLLMLoggingBaseClass):
         )
         self.function_id = function_id
         self.streaming_chunks: List[Any] = []  # for generating complete stream response
-        self.sync_streaming_chunks: List[
-            Any
-        ] = []  # for generating complete stream response
+        self.sync_streaming_chunks: List[Any] = (
+            []
+        )  # for generating complete stream response
         self.log_raw_request_response = log_raw_request_response
 
         # Initialize dynamic callbacks
@@ -746,9 +746,9 @@ class Logging(LiteLLMLoggingBaseClass):
                         prompt_spec=prompt_spec,
                         dynamic_callback_params=dynamic_callback_params,
                     ):
-                        self.model_call_details[
-                            "prompt_integration"
-                        ] = logger.__class__.__name__
+                        self.model_call_details["prompt_integration"] = (
+                            logger.__class__.__name__
+                        )
                         return logger
                 except Exception:
                     # If check fails, continue to next logger
@@ -816,9 +816,9 @@ class Logging(LiteLLMLoggingBaseClass):
         if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook(
             non_default_params
         ):
-            self.model_call_details[
-                "prompt_integration"
-            ] = anthropic_cache_control_logger.__class__.__name__
+            self.model_call_details["prompt_integration"] = (
+                anthropic_cache_control_logger.__class__.__name__
+            )
             return anthropic_cache_control_logger
 
         #########################################################
@@ -830,9 +830,9 @@ class Logging(LiteLLMLoggingBaseClass):
                 internal_usage_cache=None,
                 llm_router=None,
             )
-            self.model_call_details[
-                "prompt_integration"
-            ] = vector_store_custom_logger.__class__.__name__
+            self.model_call_details["prompt_integration"] = (
+                vector_store_custom_logger.__class__.__name__
+            )
             # Add to global callbacks so post-call hooks are invoked
             if (
                 vector_store_custom_logger
@@ -892,9 +892,9 @@ class Logging(LiteLLMLoggingBaseClass):
             model
         ):  # if model name was changes pre-call, overwrite the initial model call name with the new one
             self.model_call_details["model"] = model
-        self.model_call_details["litellm_params"][
-            "api_base"
-        ] = self._get_masked_api_base(additional_args.get("api_base", ""))
+        self.model_call_details["litellm_params"]["api_base"] = (
+            self._get_masked_api_base(additional_args.get("api_base", ""))
+        )
 
     def pre_call(self, input, api_key, model=None, additional_args={}):  # noqa: PLR0915
         # Log the exact input to the LLM API
@@ -923,10 +923,10 @@ class Logging(LiteLLMLoggingBaseClass):
                 try:
                     # [Non-blocking Extra Debug Information in metadata]
                     if turn_off_message_logging is True:
-                        _metadata[
-                            "raw_request"
-                        ] = "redacted by litellm. \
+                        _metadata["raw_request"] = (
+                            "redacted by litellm. \
                             'litellm.turn_off_message_logging=True'"
+                        )
                     else:
                         curl_command = self._get_request_curl_command(
                             api_base=additional_args.get("api_base", ""),
@@ -937,34 +937,34 @@ class Logging(LiteLLMLoggingBaseClass):
 
                         _metadata["raw_request"] = str(curl_command)
                         # split up, so it's easier to parse in the UI
-                        self.model_call_details[
-                            "raw_request_typed_dict"
-                        ] = RawRequestTypedDict(
-                            raw_request_api_base=str(
-                                additional_args.get("api_base") or ""
-                            ),
-                            raw_request_body=self._get_raw_request_body(
-                                additional_args.get("complete_input_dict", {})
-                            ),
-                            # NOTE: setting ignore_sensitive_headers to True will cause
-                            # the Authorization header to be leaked when calls to the health
-                            # endpoint are made and fail.
-                            raw_request_headers=self._get_masked_headers(
-                                additional_args.get("headers", {}) or {},
-                            ),
-                            error=None,
+                        self.model_call_details["raw_request_typed_dict"] = (
+                            RawRequestTypedDict(
+                                raw_request_api_base=str(
+                                    additional_args.get("api_base") or ""
+                                ),
+                                raw_request_body=self._get_raw_request_body(
+                                    additional_args.get("complete_input_dict", {})
+                                ),
+                                # NOTE: setting ignore_sensitive_headers to True will cause
+                                # the Authorization header to be leaked when calls to the health
+                                # endpoint are made and fail.
+                                raw_request_headers=self._get_masked_headers(
+                                    additional_args.get("headers", {}) or {},
+                                ),
+                                error=None,
+                            )
                         )
                 except Exception as e:
-                    self.model_call_details[
-                        "raw_request_typed_dict"
-                    ] = RawRequestTypedDict(
-                        error=str(e),
+                    self.model_call_details["raw_request_typed_dict"] = (
+                        RawRequestTypedDict(
+                            error=str(e),
+                        )
                     )
-                    _metadata[
-                        "raw_request"
-                    ] = "Unable to Log \
+                    _metadata["raw_request"] = (
+                        "Unable to Log \
                         raw request: {}".format(
-                        str(e)
+                            str(e)
+                        )
                     )
             if getattr(self, "logger_fn", None) and callable(self.logger_fn):
                 try:
@@ -1265,13 +1265,13 @@ class Logging(LiteLLMLoggingBaseClass):
         for callback in callbacks:
             try:
                 if isinstance(callback, CustomLogger):
-                    response: Optional[
-                        MCPPostCallResponseObject
-                    ] = await callback.async_post_mcp_tool_call_hook(
-                        kwargs=kwargs,
-                        response_obj=post_mcp_tool_call_response_obj,
-                        start_time=start_time,
-                        end_time=end_time,
+                    response: Optional[MCPPostCallResponseObject] = (
+                        await callback.async_post_mcp_tool_call_hook(
+                            kwargs=kwargs,
+                            response_obj=post_mcp_tool_call_response_obj,
+                            start_time=start_time,
+                            end_time=end_time,
+                        )
                     )
                     ######################################################################
                     # if any of the callbacks modify the response, use the modified response
@@ -1466,9 +1466,9 @@ class Logging(LiteLLMLoggingBaseClass):
             verbose_logger.debug(
                 f"response_cost_failure_debug_information: {debug_info}"
             )
-            self.model_call_details[
-                "response_cost_failure_debug_information"
-            ] = debug_info
+            self.model_call_details["response_cost_failure_debug_information"] = (
+                debug_info
+            )
             return None
 
         try:
@@ -1494,9 +1494,9 @@ class Logging(LiteLLMLoggingBaseClass):
             verbose_logger.debug(
                 f"response_cost_failure_debug_information: {debug_info}"
             )
-            self.model_call_details[
-                "response_cost_failure_debug_information"
-            ] = debug_info
+            self.model_call_details["response_cost_failure_debug_information"] = (
+                debug_info
+            )
 
         return None
 
@@ -1652,10 +1652,8 @@ class Logging(LiteLLMLoggingBaseClass):
                 result=logging_result
             )
 
-        self.model_call_details[
-            "standard_logging_object"
-        ] = self._build_standard_logging_payload(
-            logging_result, start_time, end_time
+        self.model_call_details["standard_logging_object"] = (
+            self._build_standard_logging_payload(logging_result, start_time, end_time)
         )
 
         if (
@@ -1734,9 +1732,9 @@ class Logging(LiteLLMLoggingBaseClass):
                 end_time = datetime.datetime.now()
             if self.completion_start_time is None:
                 self.completion_start_time = end_time
-                self.model_call_details[
-                    "completion_start_time"
-                ] = self.completion_start_time
+                self.model_call_details["completion_start_time"] = (
+                    self.completion_start_time
+                )
 
             self.model_call_details["log_event_type"] = "successful_api_call"
             self.model_call_details["end_time"] = end_time
@@ -1773,10 +1771,10 @@ class Logging(LiteLLMLoggingBaseClass):
                         end_time=end_time,
                     )
                 elif isinstance(result, dict) or isinstance(result, list):
-                    self.model_call_details[
-                        "standard_logging_object"
-                    ] = self._build_standard_logging_payload(
-                        result, start_time, end_time
+                    self.model_call_details["standard_logging_object"] = (
+                        self._build_standard_logging_payload(
+                            result, start_time, end_time
+                        )
                     )
                     if (
                         standard_logging_payload := self.model_call_details.get(
@@ -1785,9 +1783,9 @@ class Logging(LiteLLMLoggingBaseClass):
                     ) is not None:
                         emit_standard_logging_payload(standard_logging_payload)
             elif standard_logging_object is not None:
-                self.model_call_details[
-                    "standard_logging_object"
-                ] = standard_logging_object
+                self.model_call_details["standard_logging_object"] = (
+                    standard_logging_object
+                )
             else:
                 self.model_call_details["response_cost"] = None
 
@@ -1945,17 +1943,17 @@ class Logging(LiteLLMLoggingBaseClass):
                 verbose_logger.debug(
                     "Logging Details LiteLLM-Success Call streaming complete"
                 )
-                self.model_call_details[
-                    "complete_streaming_response"
-                ] = complete_streaming_response
-                self.model_call_details[
-                    "response_cost"
-                ] = self._response_cost_calculator(result=complete_streaming_response)
+                self.model_call_details["complete_streaming_response"] = (
+                    complete_streaming_response
+                )
+                self.model_call_details["response_cost"] = (
+                    self._response_cost_calculator(result=complete_streaming_response)
+                )
                 ## STANDARDIZED LOGGING PAYLOAD
-                self.model_call_details[
-                    "standard_logging_object"
-                ] = self._build_standard_logging_payload(
-                    complete_streaming_response, start_time, end_time
+                self.model_call_details["standard_logging_object"] = (
+                    self._build_standard_logging_payload(
+                        complete_streaming_response, start_time, end_time
+                    )
                 )
                 if (
                     standard_logging_payload := self.model_call_details.get(
@@ -2289,10 +2287,10 @@ class Logging(LiteLLMLoggingBaseClass):
                             )
                         else:
                             if self.stream and complete_streaming_response:
-                                self.model_call_details[
-                                    "complete_response"
-                                ] = self.model_call_details.get(
-                                    "complete_streaming_response", {}
+                                self.model_call_details["complete_response"] = (
+                                    self.model_call_details.get(
+                                        "complete_streaming_response", {}
+                                    )
                                 )
                                 result = self.model_call_details["complete_response"]
                             openMeterLogger.log_success_event(
@@ -2316,10 +2314,10 @@ class Logging(LiteLLMLoggingBaseClass):
                             )
                         else:
                             if self.stream and complete_streaming_response:
-                                self.model_call_details[
-                                    "complete_response"
-                                ] = self.model_call_details.get(
-                                    "complete_streaming_response", {}
+                                self.model_call_details["complete_response"] = (
+                                    self.model_call_details.get(
+                                        "complete_streaming_response", {}
+                                    )
                                 )
                                 result = self.model_call_details["complete_response"]
 
@@ -2458,9 +2456,9 @@ class Logging(LiteLLMLoggingBaseClass):
         if complete_streaming_response is not None:
             print_verbose("Async success callbacks: Got a complete streaming response")
 
-            self.model_call_details[
-                "async_complete_streaming_response"
-            ] = complete_streaming_response
+            self.model_call_details["async_complete_streaming_response"] = (
+                complete_streaming_response
+            )
 
             try:
                 if self.model_call_details.get("cache_hit", False) is True:
@@ -2471,10 +2469,10 @@ class Logging(LiteLLMLoggingBaseClass):
                         model_call_details=self.model_call_details
                     )
                     # base_model defaults to None if not set on model_info
-                    self.model_call_details[
-                        "response_cost"
-                    ] = self._response_cost_calculator(
-                        result=complete_streaming_response
+                    self.model_call_details["response_cost"] = (
+                        self._response_cost_calculator(
+                            result=complete_streaming_response
+                        )
                     )
 
                 verbose_logger.debug(
@@ -2487,10 +2485,10 @@ class Logging(LiteLLMLoggingBaseClass):
                 self.model_call_details["response_cost"] = None
 
             ## STANDARDIZED LOGGING PAYLOAD
-            self.model_call_details[
-                "standard_logging_object"
-            ] = self._build_standard_logging_payload(
-                complete_streaming_response, start_time, end_time
+            self.model_call_details["standard_logging_object"] = (
+                self._build_standard_logging_payload(
+                    complete_streaming_response, start_time, end_time
+                )
             )
 
             # print standard logging payload
@@ -2517,10 +2515,8 @@ class Logging(LiteLLMLoggingBaseClass):
             # _success_handler_helper_fn
             if self.model_call_details.get("standard_logging_object") is None:
                 ## STANDARDIZED LOGGING PAYLOAD
-                self.model_call_details[
-                    "standard_logging_object"
-                ] = self._build_standard_logging_payload(
-                    result, start_time, end_time
+                self.model_call_details["standard_logging_object"] = (
+                    self._build_standard_logging_payload(result, start_time, end_time)
                 )
 
                 # print standard logging payload
@@ -2764,18 +2760,18 @@ class Logging(LiteLLMLoggingBaseClass):
 
         ## STANDARDIZED LOGGING PAYLOAD
 
-        self.model_call_details[
-            "standard_logging_object"
-        ] = get_standard_logging_object_payload(
-            kwargs=self.model_call_details,
-            init_response_obj={},
-            start_time=start_time,
-            end_time=end_time,
-            logging_obj=self,
-            status="failure",
-            error_str=str(exception),
-            original_exception=exception,
-            standard_built_in_tools_params=self.standard_built_in_tools_params,
+        self.model_call_details["standard_logging_object"] = (
+            get_standard_logging_object_payload(
+                kwargs=self.model_call_details,
+                init_response_obj={},
+                start_time=start_time,
+                end_time=end_time,
+                logging_obj=self,
+                status="failure",
+                error_str=str(exception),
+                original_exception=exception,
+                standard_built_in_tools_params=self.standard_built_in_tools_params,
+            )
         )
         return start_time, end_time
 
@@ -3739,9 +3735,9 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
                 service_name=arize_config.project_name,
             )
 
-            os.environ[
-                "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
-            ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}"
+            os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
+                f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}"
+            )
             for callback in _in_memory_loggers:
                 if (
                     isinstance(callback, ArizeLogger)
@@ -3767,13 +3763,13 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
                 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "")
                 # Add openinference.project.name attribute
                 if existing_attrs:
-                    os.environ[
-                        "OTEL_RESOURCE_ATTRIBUTES"
-                    ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}"
+                    os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
+                        f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}"
+                    )
                 else:
-                    os.environ[
-                        "OTEL_RESOURCE_ATTRIBUTES"
-                    ] = f"openinference.project.name={arize_phoenix_config.project_name}"
+                    os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
+                        f"openinference.project.name={arize_phoenix_config.project_name}"
+                    )
 
             # Set Phoenix project name from environment variable
             phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None)
@@ -3781,19 +3777,19 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
                 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "")
                 # Add openinference.project.name attribute
                 if existing_attrs:
-                    os.environ[
-                        "OTEL_RESOURCE_ATTRIBUTES"
-                    ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}"
+                    os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
+                        f"{existing_attrs},openinference.project.name={phoenix_project_name}"
+                    )
                 else:
-                    os.environ[
-                        "OTEL_RESOURCE_ATTRIBUTES"
-                    ] = f"openinference.project.name={phoenix_project_name}"
+                    os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
+                        f"openinference.project.name={phoenix_project_name}"
+                    )
 
             # auth can be disabled on local deployments of arize phoenix
             if arize_phoenix_config.otlp_auth_headers is not None:
-                os.environ[
-                    "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
-                ] = arize_phoenix_config.otlp_auth_headers
+                os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
+                    arize_phoenix_config.otlp_auth_headers
+                )
 
             for callback in _in_memory_loggers:
                 if (
@@ -3969,9 +3965,9 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
                 exporter="otlp_http",
                 endpoint="https://langtrace.ai/api/trace",
             )
-            os.environ[
-                "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
-            ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}"
+            os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
+                f"api_key={os.getenv('LANGTRACE_API_KEY')}"
+            )
             for callback in _in_memory_loggers:
                 if (
                     isinstance(callback, OpenTelemetry)
@@ -4204,8 +4200,7 @@ def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None:
         litellm.logging_callback_manager.add_litellm_callback(phoenix_logger)
 
         verbose_logger.info(
-            "Auto-initialized Arize Phoenix logger alongside otel "
-            "(endpoint=%s)",
+            "Auto-initialized Arize Phoenix logger alongside otel " "(endpoint=%s)",
             arize_phoenix_config.endpoint,
         )
     except Exception as e:
@@ -4768,9 +4763,11 @@ class StandardLoggingPayloadSetup:
             ).model_dump()
         if isinstance(_raw, dict):
             if ResponseAPILoggingUtils._is_response_api_usage(_raw):
-                return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
-                    _raw
-                ).model_dump()
+                return (
+                    ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
+                        _raw
+                    ).model_dump()
+                )
             return _raw
         if isinstance(_raw, Usage):
             return _raw.model_dump()
@@ -4884,10 +4881,10 @@ class StandardLoggingPayloadSetup:
             for key in StandardLoggingHiddenParams.__annotations__.keys():
                 if key in hidden_params:
                     if key == "additional_headers":
-                        clean_hidden_params[
-                            "additional_headers"
-                        ] = StandardLoggingPayloadSetup.get_additional_headers(
-                            hidden_params[key]
+                        clean_hidden_params["additional_headers"] = (
+                            StandardLoggingPayloadSetup.get_additional_headers(
+                                hidden_params[key]
+                            )
                         )
                     else:
                         clean_hidden_params[key] = hidden_params[key]  # type: ignore
@@ -5039,14 +5036,22 @@ class StandardLoggingPayloadSetup:
         dynamic_litellm_session_id = litellm_params.get("litellm_session_id")
         dynamic_litellm_trace_id = litellm_params.get("litellm_trace_id")
 
+
         # Note: we recommend using `litellm_session_id` for session tracking
         # `litellm_trace_id` is an internal litellm param
         if dynamic_litellm_session_id:
             return str(dynamic_litellm_session_id)
         elif dynamic_litellm_trace_id:
             return str(dynamic_litellm_trace_id)
-        else:
-            return logging_obj.litellm_trace_id
+        # Fallback: use metadata.session_id or metadata.trace_id for call chaining
+        metadata = litellm_params.get("metadata") or {}
+        metadata_session_id = metadata.get("session_id")
+        metadata_trace_id = metadata.get("trace_id")
+        if metadata_session_id:
+            return str(metadata_session_id)
+        if metadata_trace_id:
+            return str(metadata_trace_id)
+        return logging_obj.litellm_trace_id
 
     @staticmethod
     def _get_user_agent_tags(proxy_server_request: dict) -> Optional[List[str]]:
@@ -5502,9 +5507,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]):
     ):
         for k, v in metadata["user_api_key_metadata"].items():
             if k == "logging":  # prevent logging user logging keys
-                cleaned_user_api_key_metadata[
-                    k
-                ] = "scrubbed_by_litellm_for_sensitive_keys"
+                cleaned_user_api_key_metadata[k] = (
+                    "scrubbed_by_litellm_for_sensitive_keys"
+                )
             else:
                 cleaned_user_api_key_metadata[k] = v
 
@@ -5616,4 +5621,3 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
         model_parameters={"stream": True},
         hidden_params=hidden_params,
     )
-
diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 5b3d5bd60e2..cdb26acb658 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -5,7 +5,6 @@ LiteLLM MCP Server Routes
 
 import asyncio
 import contextlib
-
 import traceback
 import uuid
 from datetime import datetime
@@ -44,7 +43,10 @@ from litellm.proxy._experimental.mcp_server.utils import (
 )
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.ip_address_utils import IPAddressUtils
-from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
+from litellm.proxy.litellm_pre_call_utils import (
+    LiteLLMProxyRequestSetup,
+    get_chain_id_from_headers,
+)
 from litellm.types.mcp import MCPAuth
 from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
 from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall
@@ -331,6 +333,11 @@ if MCP_AVAILABLE:
         try:
             # Create a body date for logging
             body_data = {"name": name, "arguments": arguments}
+            # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A)
+            chain_id = get_chain_id_from_headers(raw_headers)
+            if chain_id:
+                body_data["litellm_trace_id"] = chain_id
+                body_data["litellm_session_id"] = chain_id
 
             request = Request(
                 scope={
@@ -884,6 +891,10 @@ if MCP_AVAILABLE:
             # This is intentionally minimal: only async_success_handler / post_call_failure_hook
             rules_obj = Rules()
             list_tools_call_id = str(uuid.uuid4())
+            # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool)
+            effective_litellm_trace_id = litellm_trace_id or get_chain_id_from_headers(
+                raw_headers
+            )
             spend_logs_metadata: Dict[str, Any] = {
                 "mcp_operation": "list_tools",
             }
@@ -896,7 +907,7 @@ if MCP_AVAILABLE:
                 "model": "MCP: list_tools",
                 "call_type": CallTypes.list_mcp_tools.value,
                 "litellm_call_id": list_tools_call_id,
-                "litellm_trace_id": litellm_trace_id,
+                "litellm_trace_id": effective_litellm_trace_id,
                 "metadata": {
                     "spend_logs_metadata": spend_logs_metadata,
                 },
diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/404.html
rename to litellm/proxy/_experimental/out/404/index.html
diff --git a/litellm/proxy/_experimental/out/_not-found.html b/litellm/proxy/_experimental/out/_not-found/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/_not-found.html
rename to litellm/proxy/_experimental/out/_not-found/index.html
diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/api-reference.html
rename to litellm/proxy/_experimental/out/api-reference/index.html
diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/experimental/api-playground.html
rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html
diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/experimental/budgets.html
rename to litellm/proxy/_experimental/out/experimental/budgets/index.html
diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/experimental/caching.html
rename to litellm/proxy/_experimental/out/experimental/caching/index.html
diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html
rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html
diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/experimental/old-usage.html
rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html
diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/experimental/prompts.html
rename to litellm/proxy/_experimental/out/experimental/prompts/index.html
diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/experimental/tag-management.html
rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html
diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/guardrails.html
rename to litellm/proxy/_experimental/out/guardrails/index.html
diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/login.html
rename to litellm/proxy/_experimental/out/login/index.html
diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/logs.html
rename to litellm/proxy/_experimental/out/logs/index.html
diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html
rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html
diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/model-hub.html
rename to litellm/proxy/_experimental/out/model-hub/index.html
diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/model_hub.html
rename to litellm/proxy/_experimental/out/model_hub/index.html
diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/model_hub_table.html
rename to litellm/proxy/_experimental/out/model_hub_table/index.html
diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/models-and-endpoints.html
rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html
diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/onboarding.html
rename to litellm/proxy/_experimental/out/onboarding/index.html
diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/organizations.html
rename to litellm/proxy/_experimental/out/organizations/index.html
diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/playground.html
rename to litellm/proxy/_experimental/out/playground/index.html
diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/policies.html
rename to litellm/proxy/_experimental/out/policies/index.html
diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/settings/admin-settings.html
rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html
diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html
rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html
diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/settings/router-settings.html
rename to litellm/proxy/_experimental/out/settings/router-settings/index.html
diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/settings/ui-theme.html
rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html
diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/teams.html
rename to litellm/proxy/_experimental/out/teams/index.html
diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/test-key.html
rename to litellm/proxy/_experimental/out/test-key/index.html
diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/tools/mcp-servers.html
rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html
diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/tools/vector-stores.html
rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html
diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/usage.html
rename to litellm/proxy/_experimental/out/usage/index.html
diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/users.html
rename to litellm/proxy/_experimental/out/users/index.html
diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html
similarity index 100%
rename from litellm/proxy/_experimental/out/virtual-keys.html
rename to litellm/proxy/_experimental/out/virtual-keys/index.html
diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py
index 7f30277ebca..6bcee14f29e 100644
--- a/litellm/proxy/agent_endpoints/a2a_endpoints.py
+++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py
@@ -69,6 +69,7 @@ async def _handle_stream_message(
     from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE
 
     if not A2A_SDK_AVAILABLE:
+
         async def _error_stream():
             yield json.dumps(
                 {
@@ -106,7 +107,12 @@ async def _handle_stream_message(
                 proxy_server_request=proxy_server_request,
             )
 
-            if use_proxy_hooks and user_api_key_dict is not None and request_data is not None and proxy_logging_obj is not None:
+            if (
+                use_proxy_hooks
+                and user_api_key_dict is not None
+                and request_data is not None
+                and proxy_logging_obj is not None
+            ):
                 from litellm.proxy.common_request_processing import (
                     ProxyBaseLLMRequestProcessing,
                 )
@@ -119,20 +125,27 @@ async def _handle_stream_message(
                     return json.dumps(obj) + "\n"
 
                 def _ndjson_error(proxy_exc: Any) -> str:
-                    return json.dumps(
-                        {
-                            "jsonrpc": "2.0",
-                            "id": request_id,
-                            "error": {
-                                "code": -32603,
-                                "message": getattr(
-                                    proxy_exc, "message", f"Streaming error: {proxy_exc!s}"
-                                ),
-                            },
-                        }
-                    ) + "\n"
+                    return (
+                        json.dumps(
+                            {
+                                "jsonrpc": "2.0",
+                                "id": request_id,
+                                "error": {
+                                    "code": -32603,
+                                    "message": getattr(
+                                        proxy_exc,
+                                        "message",
+                                        f"Streaming error: {proxy_exc!s}",
+                                    ),
+                                },
+                            }
+                        )
+                        + "\n"
+                    )
 
-                async for line in ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
+                async for (
+                    line
+                ) in ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
                     response=a2a_stream,
                     user_api_key_dict=user_api_key_dict,
                     request_data=request_data,
@@ -151,7 +164,12 @@ async def _handle_stream_message(
                         yield json.dumps(chunk) + "\n"
         except Exception as e:
             verbose_proxy_logger.exception(f"Error streaming A2A response: {e}")
-            if use_proxy_hooks and proxy_logging_obj is not None and user_api_key_dict is not None and request_data is not None:
+            if (
+                use_proxy_hooks
+                and proxy_logging_obj is not None
+                and user_api_key_dict is not None
+                and request_data is not None
+            ):
                 transformed_exception = await proxy_logging_obj.post_call_failure_hook(
                     user_api_key_dict=user_api_key_dict,
                     original_exception=e,
@@ -382,6 +400,7 @@ async def invoke_agent_a2a(
                 agent_id=agent.agent_id,
                 metadata=data.get("metadata", {}),
                 proxy_server_request=data.get("proxy_server_request"),
+                litellm_logging_obj=logging_obj,
             )
 
             response = await proxy_logging_obj.post_call_success_hook(
diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py
index c6a709534e1..4c866a24991 100644
--- a/litellm/proxy/guardrails/guardrail_endpoints.py
+++ b/litellm/proxy/guardrails/guardrail_endpoints.py
@@ -1624,11 +1624,11 @@ def _build_field_dict(
     # Determine the field type from annotation
     field_type = _get_field_type_from_annotation(field_annotation)
 
-    # Check for custom UI type override (ui_type preferred; "type" leaks into OpenAPI and breaks schema)
-    field_json_schema_extra = getattr(field, "json_schema_extra", {}) or {}
+    # Check for custom UI type override
+    field_json_schema_extra = getattr(field, "json_schema_extra", {})
     if field_json_schema_extra and "ui_type" in field_json_schema_extra:
-        ut = field_json_schema_extra["ui_type"]
-        field_type = ut if isinstance(ut, str) else getattr(ut, "value", ut)
+        ui_type = field_json_schema_extra["ui_type"]
+        field_type = ui_type.value if hasattr(ui_type, "value") else ui_type
     elif field_json_schema_extra and "type" in field_json_schema_extra:
         field_type = field_json_schema_extra["type"]
 
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index 92bf035a986..68f1ab114bf 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -89,6 +89,25 @@ def _get_metadata_variable_name(request: Request) -> str:
     return "metadata"
 
 
+def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str]:
+    """
+    Extract chain id for call chaining from request headers.
+
+    x-litellm-trace-id and x-litellm-session-id are interchangeable; when both
+    are present, x-litellm-trace-id takes precedence. Header keys are matched
+    case-insensitively so this works with raw header dicts from any transport.
+
+    Used by MCP (and other paths that have raw_headers but no Request) to set
+    litellm_trace_id/litellm_session_id for spend logs and logging consistency.
+    """
+    if not headers:
+        return None
+    normalized = {k.lower(): v for k, v in headers.items() if isinstance(k, str)}
+    return normalized.get("x-litellm-trace-id") or normalized.get(
+        "x-litellm-session-id"
+    )
+
+
 def safe_add_api_version_from_query_params(data: dict, request: Request):
     try:
         if hasattr(request, "query_params"):
@@ -177,12 +196,12 @@ def _get_dynamic_logging_metadata(
     user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig
 ) -> Optional[TeamCallbackMetadata]:
     callback_settings_obj: Optional[TeamCallbackMetadata] = None
-    key_dynamic_logging_settings: Optional[
-        dict
-    ] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
-    team_dynamic_logging_settings: Optional[
-        dict
-    ] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
+    key_dynamic_logging_settings: Optional[dict] = (
+        KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
+    )
+    team_dynamic_logging_settings: Optional[dict] = (
+        KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
+    )
     #########################################################################################
     # Key-based callbacks
     #########################################################################################
@@ -576,9 +595,13 @@ class LiteLLMProxyRequestSetup:
         #########################################################################################
         # Finally update the requests metadata with the `metadata_from_headers`
         #########################################################################################
+
         agent_id_from_header = headers.get("x-litellm-agent-id")
-        trace_id_from_header = headers.get("x-litellm-trace-id")
-        session_id_from_header = headers.get("x-litellm-session-id")
+        # x-litellm-trace-id and x-litellm-session-id are interchangeable for call chaining
+        chain_id = headers.get("x-litellm-trace-id") or headers.get(
+            "x-litellm-session-id"
+        )
+
 
         if agent_id_from_header:
             metadata_from_headers["agent_id"] = agent_id_from_header
@@ -586,16 +609,13 @@ class LiteLLMProxyRequestSetup:
                 f"Extracted agent_id from header: {agent_id_from_header}"
             )
 
-        if trace_id_from_header:
-            metadata_from_headers["trace_id"] = trace_id_from_header
+        if chain_id:
+            metadata_from_headers["trace_id"] = chain_id
+            metadata_from_headers["session_id"] = chain_id
+            data["litellm_session_id"] = chain_id
+            data["litellm_trace_id"] = chain_id
             verbose_proxy_logger.debug(
-                f"Extracted trace_id from header: {trace_id_from_header}"
-            )
-
-        if session_id_from_header:
-            metadata_from_headers["session_id"] = session_id_from_header
-            verbose_proxy_logger.debug(
-                f"Extracted session_id from header: {session_id_from_header}"
+                f"Extracted chain_id from header (trace-id/session-id): {chain_id}"
             )
 
         if isinstance(data[_metadata_variable_name], dict):
@@ -702,11 +722,11 @@ class LiteLLMProxyRequestSetup:
 
         ## KEY-LEVEL SPEND LOGS / TAGS
         if "tags" in key_metadata and key_metadata["tags"] is not None:
-            data[_metadata_variable_name][
-                "tags"
-            ] = LiteLLMProxyRequestSetup._merge_tags(
-                request_tags=data[_metadata_variable_name].get("tags"),
-                tags_to_add=key_metadata["tags"],
+            data[_metadata_variable_name]["tags"] = (
+                LiteLLMProxyRequestSetup._merge_tags(
+                    request_tags=data[_metadata_variable_name].get("tags"),
+                    tags_to_add=key_metadata["tags"],
+                )
             )
         if "disable_global_guardrails" in key_metadata and isinstance(
             key_metadata["disable_global_guardrails"], bool
@@ -839,14 +859,9 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
     """
 
     from litellm.proxy.proxy_server import llm_router, premium_user
-    from litellm.types.proxy.litellm_pre_call_utils import (
-        RedactedDict,
-        SecretFields,
-    )
+    from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields
 
-    _raw_headers: Dict[str, str] = RedactedDict(
-        _safe_get_request_headers(request)
-    )
+    _raw_headers: Dict[str, str] = RedactedDict(_safe_get_request_headers(request))
 
     forward_llm_auth = False
     if general_settings:
@@ -986,9 +1001,9 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
     data[_metadata_variable_name]["litellm_api_version"] = version
 
     if general_settings is not None:
-        data[_metadata_variable_name][
-            "global_max_parallel_requests"
-        ] = general_settings.get("global_max_parallel_requests", None)
+        data[_metadata_variable_name]["global_max_parallel_requests"] = (
+            general_settings.get("global_max_parallel_requests", None)
+        )
 
     ### KEY-LEVEL Controls
     key_metadata = user_api_key_dict.metadata
diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py
index 31615a768d7..131841f7b59 100644
--- a/litellm/proxy/spend_tracking/spend_tracking_utils.py
+++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py
@@ -11,26 +11,21 @@ from pydantic import BaseModel
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.constants import (
-    MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB,
-)
+from litellm.constants import \
+    MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB
 from litellm.constants import REDACTED_BY_LITELM_STRING
 from litellm.litellm_core_utils.core_helpers import (
-    get_litellm_metadata_from_kwargs,
-    reconstruct_model_name,
-)
+    get_litellm_metadata_from_kwargs, reconstruct_model_name)
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
 from litellm.proxy.utils import PrismaClient, hash_token
-from litellm.types.utils import (
-    CostBreakdown,
-    StandardLoggingGuardrailInformation,
-    StandardLoggingMCPToolCall,
-    StandardLoggingModelInformation,
-    StandardLoggingPayload,
-    StandardLoggingVectorStoreRequest,
-    VectorStoreSearchResponse,
-)
+from litellm.types.utils import (CostBreakdown,
+                                 StandardLoggingGuardrailInformation,
+                                 StandardLoggingMCPToolCall,
+                                 StandardLoggingModelInformation,
+                                 StandardLoggingPayload,
+                                 StandardLoggingVectorStoreRequest,
+                                 VectorStoreSearchResponse)
 from litellm.utils import get_end_user_id_for_cost_tracking
 
 
@@ -116,16 +111,15 @@ def _get_spend_logs_metadata(
     # Filter the metadata dictionary to include only the specified keys
     clean_metadata = SpendLogsMetadata(
         **{  # type: ignore
-            key: metadata.get(key)
-            for key in SpendLogsMetadata.__annotations__.keys()
+            key: metadata.get(key) for key in SpendLogsMetadata.__annotations__.keys()
         }
     )
     clean_metadata["applied_guardrails"] = applied_guardrails
     clean_metadata["batch_models"] = batch_models
     clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata
-    clean_metadata[
-        "vector_store_request_metadata"
-    ] = _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata)
+    clean_metadata["vector_store_request_metadata"] = (
+        _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata)
+    )
     clean_metadata["guardrail_information"] = guardrail_information
     clean_metadata["usage_object"] = usage_object
     clean_metadata["model_map_information"] = model_map_information
@@ -372,9 +366,11 @@ def get_logging_payload(  # noqa: PLR0915
         guardrail_information=(
             standard_logging_payload.get("guardrail_information", None)
             if standard_logging_payload is not None
-            else metadata.get("standard_logging_guardrail_information", None)
-            if metadata is not None
-            else None
+            else (
+                metadata.get("standard_logging_guardrail_information", None)
+                if metadata is not None
+                else None
+            )
         ),
         cold_storage_object_key=(
             standard_logging_payload["metadata"].get("cold_storage_object_key", None)
@@ -501,6 +497,7 @@ def _get_session_id_for_spend_log(
     """
     from litellm._uuid import uuid
 
+
     if (
         standard_logging_payload is not None
         and standard_logging_payload.get("trace_id") is not None
@@ -515,9 +512,7 @@ def _get_session_id_for_spend_log(
     return str(uuid.uuid4())
 
 
-def _get_request_duration_ms(
-    start_time: datetime, end_time: datetime
-) -> Optional[int]:
+def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> Optional[int]:
     """Compute request duration in milliseconds from start and end times."""
     try:
         return int((end_time - start_time).total_seconds() * 1000)
@@ -709,20 +704,20 @@ def _convert_to_json_serializable_dict(
     if max_depth <= 0:
         # Return a placeholder if max depth is exceeded
         return ""
-    
+
     if visited is None:
         visited = set()
-    
+
     # Get the object's memory address to track visited objects
     obj_id = id(obj)
     if obj_id in visited:
         # Circular reference detected, return placeholder
         return ""
-    
+
     # Only track mutable objects (dict, list, objects with __dict__)
     if isinstance(obj, (dict, list)) or hasattr(obj, "__dict__"):
         visited.add(obj_id)
-    
+
     try:
         if isinstance(obj, BaseModel):
             # Use Pydantic's model_dump() instead of pickle
@@ -741,7 +736,9 @@ def _convert_to_json_serializable_dict(
             ]
         elif hasattr(obj, "__dict__"):
             # Handle objects with __dict__ attribute
-            return _convert_to_json_serializable_dict(obj.__dict__, visited, max_depth - 1)
+            return _convert_to_json_serializable_dict(
+                obj.__dict__, visited, max_depth - 1
+            )
         else:
             # Primitives (str, int, float, bool, None) pass through
             return obj
@@ -777,9 +774,7 @@ def _get_proxy_server_request_for_spend_logs_payload(
             # Apply message redaction if turn_off_message_logging is enabled
             if kwargs is not None:
                 from litellm.litellm_core_utils.redact_messages import (
-                    perform_redaction,
-                    should_redact_message_logging,
-                )
+                    perform_redaction, should_redact_message_logging)
 
                 # Build model_call_details dict to check redaction settings
                 model_call_details = {
@@ -788,12 +783,12 @@ def _get_proxy_server_request_for_spend_logs_payload(
                         "standard_callback_dynamic_params"
                     ),
                 }
-                
+
                 # If redaction is enabled, convert to serializable dict before redacting
                 if should_redact_message_logging(model_call_details=model_call_details):
                     _request_body = _convert_to_json_serializable_dict(_request_body)
                     perform_redaction(model_call_details=_request_body, result=None)
-            
+
             _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body)
             _request_body_json_str = json.dumps(_request_body, default=str)
             return _request_body_json_str
@@ -845,10 +840,8 @@ def _get_response_for_spend_logs_payload(
         # Apply message redaction if turn_off_message_logging is enabled
         if kwargs is not None:
             from litellm.litellm_core_utils.redact_messages import (
-                perform_redaction,
-                should_redact_message_logging,
-            )
-            
+                perform_redaction, should_redact_message_logging)
+
             litellm_params = kwargs.get("litellm_params", {})
             model_call_details = {
                 "litellm_params": litellm_params,
@@ -856,11 +849,13 @@ def _get_response_for_spend_logs_payload(
                     "standard_callback_dynamic_params"
                 ),
             }
-            
+
             # If redaction is enabled, convert to serializable dict before redacting
             if should_redact_message_logging(model_call_details=model_call_details):
                 response_obj = _convert_to_json_serializable_dict(response_obj)
-                response_obj = perform_redaction(model_call_details={}, result=response_obj)
+                response_obj = perform_redaction(
+                    model_call_details={}, result=response_obj
+                )
 
         sanitized_wrapper = _sanitize_request_body_for_spend_logs_payload(
             {"response": response_obj}
@@ -882,7 +877,7 @@ def _should_store_prompts_and_responses_in_spend_logs() -> bool:
 
     # Check general_settings (from DB or proxy_config.yaml)
     store_prompts_value = general_settings.get("store_prompts_in_spend_logs")
-    
+
     # Normalize case: handle True/true/TRUE, False/false/FALSE, None/null
     if store_prompts_value is True:
         return True
@@ -890,7 +885,7 @@ def _should_store_prompts_and_responses_in_spend_logs() -> bool:
         # Case-insensitive string comparison
         if store_prompts_value.lower() == "true":
             return True
-    
+
     # Also check environment variable
     return get_secret_bool("STORE_PROMPTS_IN_SPEND_LOGS") is True
 
diff --git a/litellm/utils.py b/litellm/utils.py
index d192609eead..ad1eb7aeceb 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -1454,10 +1454,12 @@ def client(original_function):  # noqa: PLR0915
                 logging_obj, kwargs = function_setup(
                     original_function.__name__, rules_obj, start_time, *args, **kwargs
                 )
-            
+
             # Type assertion: logging_obj is guaranteed to be non-None after function_setup
-            assert logging_obj is not None, "logging_obj should not be None after function_setup"
-            
+            assert (
+                logging_obj is not None
+            ), "logging_obj should not be None after function_setup"
+
             ## LOAD CREDENTIALS
             load_credentials_from_list(kwargs)
             kwargs["litellm_logging_obj"] = logging_obj
@@ -1753,7 +1755,9 @@ def client(original_function):  # noqa: PLR0915
         print_args_passed_to_litellm(original_function, args, kwargs)
         start_time = datetime.datetime.now()
         result = None
-        _update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata")
+        _update_response_metadata = getattr(
+            sys.modules[__name__], "update_response_metadata"
+        )
         logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get(
             "litellm_logging_obj", None
         )
@@ -1776,9 +1780,11 @@ def client(original_function):  # noqa: PLR0915
                 logging_obj, kwargs = function_setup(
                     original_function.__name__, rules_obj, start_time, *args, **kwargs
                 )
-            
+
             # Type assertion: logging_obj is guaranteed to be non-None after function_setup
-            assert logging_obj is not None, "logging_obj should not be None after function_setup"
+            assert (
+                logging_obj is not None
+            ), "logging_obj should not be None after function_setup"
 
             modified_kwargs = await async_pre_call_deployment_hook(kwargs, call_type)
             if modified_kwargs is not None:
@@ -1861,6 +1867,7 @@ def client(original_function):  # noqa: PLR0915
             # MODEL CALL
             result = await original_function(*args, **kwargs)
             end_time = datetime.datetime.now()
+
             if _is_streaming_request(
                 kwargs=kwargs,
                 call_type=call_type,
@@ -2082,12 +2089,14 @@ def _is_async_request(
     return False
 
 
-_STREAMING_CALL_TYPES = frozenset({
-    CallTypes.generate_content_stream,
-    CallTypes.agenerate_content_stream,
-    CallTypes.generate_content_stream.value,
-    CallTypes.agenerate_content_stream.value,
-})
+_STREAMING_CALL_TYPES = frozenset(
+    {
+        CallTypes.generate_content_stream,
+        CallTypes.agenerate_content_stream,
+        CallTypes.generate_content_stream.value,
+        CallTypes.agenerate_content_stream.value,
+    }
+)
 
 
 def _is_streaming_request(
@@ -2181,7 +2190,7 @@ def encode(model="", text="", custom_tokenizer: Optional[dict] = None):
     # Normalize: HuggingFace Tokenizer.encode() returns an Encoding object;
     # extract .ids so the return type is always List[int].
     if hasattr(enc, "ids"):
-        return enc.ids
+        return enc.ids  # type: ignore
     return enc
 
 
@@ -5836,7 +5845,7 @@ def get_model_info(
                 _model_info[key] = value  # type: ignore
 
     # if verbose_logger.isEnabledFor(logging.DEBUG):
-        # verbose_logger.debug(f"model_info: {_model_info}")
+    # verbose_logger.debug(f"model_info: {_model_info}")
 
     returned_model_info = ModelInfo(
         **_model_info, supported_openai_params=supported_openai_params
@@ -6179,8 +6188,10 @@ def validate_environment(  # noqa: PLR0915
                 "AWS_ROLE_ARN" in os.environ
                 or "AWS_PROFILE" in os.environ
                 or "AWS_WEB_IDENTITY_TOKEN_FILE" in os.environ
-                or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" in os.environ  # ECS task role
-                or "AWS_CONTAINER_CREDENTIALS_FULL_URI" in os.environ  # ECS/Fargate full URI credential delivery
+                or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"
+                in os.environ  # ECS task role
+                or "AWS_CONTAINER_CREDENTIALS_FULL_URI"
+                in os.environ  # ECS/Fargate full URI credential delivery
             ):
                 keys_in_environment = True
             else:
@@ -7386,7 +7397,9 @@ class ModelResponseIterator:
         if convert_to_delta is True:
             _stream_response = ModelResponseStream()
             _stream_response.choices[0].delta.content = model_response.choices[0].message.content  # type: ignore
-            self.model_response: Union[ModelResponse, ModelResponseStream] = _stream_response
+            self.model_response: Union[ModelResponse, ModelResponseStream] = (
+                _stream_response
+            )
         else:
             self.model_response = model_response
         self.is_done = False
@@ -7457,13 +7470,13 @@ def is_cached_message(message: AllMessageValues) -> bool:
     Used for anthropic/gemini context caching.
 
     Follows the anthropic format {"cache_control": {"type": "ephemeral"}}
-    
+
     Can be disabled globally by setting litellm.disable_anthropic_gemini_context_caching_transform = True
     """
     # Check if context caching is disabled globally
     if litellm.disable_anthropic_gemini_context_caching_transform is True:
         return False
-    
+
     if "content" not in message:
         return False
 
@@ -7980,6 +7993,7 @@ class ProviderConfigManager:
     def _get_azure_ai_config(model: str) -> BaseConfig:
         """Get Azure AI config based on model type."""
         from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
         return AzureFoundryModelInfo.get_azure_ai_config_for_model(model)
 
     @staticmethod
diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
index 62a6e777b0d..ca224726361 100644
--- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
+++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
@@ -1112,6 +1112,47 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker):
     assert result.guardrail_definition_location == "db"
 
 
+class TestBuildFieldDict:
+    """Test _build_field_dict handles both enum and string ui_type values."""
+
+    def test_build_field_dict_with_string_ui_type(self):
+        """Test that _build_field_dict works when ui_type is a plain string (e.g. BlockCodeExecutionGuardrailConfigModel)."""
+        from unittest.mock import MagicMock
+
+        from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict
+
+        field = MagicMock()
+        field.json_schema_extra = {"ui_type": "multiselect", "options": ["python", "javascript"]}
+
+        result = _build_field_dict(
+            field=field,
+            field_annotation=str,
+            description="Test field",
+            required=False,
+        )
+
+        assert result["type"] == "multiselect"
+        assert result["description"] == "Test field"
+
+    def test_build_field_dict_with_enum_ui_type(self):
+        """Test that _build_field_dict works when ui_type is a GuardrailParamUITypes enum."""
+        from unittest.mock import MagicMock
+
+        from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict
+        from litellm.types.guardrails import GuardrailParamUITypes
+
+        field = MagicMock()
+        field.json_schema_extra = {"ui_type": GuardrailParamUITypes.BOOL}
+
+        result = _build_field_dict(
+            field=field,
+            field_annotation=bool,
+            description="Test bool field",
+            required=True,
+        )
+
+        assert result["type"] == "bool"
+        assert result["required"] is True
 # --- Team guardrail registration (register / submissions) ---
 
 MOCK_REGISTER_REQUEST = RegisterGuardrailRequest(
@@ -1571,4 +1612,4 @@ async def test_list_submissions_summary_counts_unaffected_by_filters(mocker):
     assert len(result.submissions) == 1  # filtered
     assert result.summary.total == 2  # unfiltered
     assert result.summary.pending_review == 1
-    assert result.summary.active == 1
\ No newline at end of file
+    assert result.summary.active == 1
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 8abc6bfe077..bc13cea939e 100644
--- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py
+++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py
@@ -11,11 +11,16 @@ from fastapi import Request
 import litellm
 from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth
 from litellm.proxy.litellm_pre_call_utils import (
-    KeyAndTeamLoggingSettings, LiteLLMProxyRequestSetup,
-    _get_dynamic_logging_metadata, _get_enforced_params,
-    _get_metadata_variable_name, _update_model_if_key_alias_exists,
-    add_guardrails_from_policy_engine, add_litellm_data_to_request,
-    check_if_token_is_service_account)
+    KeyAndTeamLoggingSettings,
+    LiteLLMProxyRequestSetup,
+    _get_dynamic_logging_metadata,
+    _get_enforced_params,
+    _get_metadata_variable_name,
+    _update_model_if_key_alias_exists,
+    add_guardrails_from_policy_engine,
+    add_litellm_data_to_request,
+    check_if_token_is_service_account,
+)
 
 sys.path.insert(
     0, os.path.abspath("../../..")
@@ -154,8 +159,7 @@ def test_get_enforced_params(
 
 @pytest.mark.asyncio
 async def test_add_litellm_data_to_request_parses_string_metadata():
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     # Setup
     request_mock = MagicMock(spec=Request)
@@ -201,8 +205,7 @@ async def test_add_litellm_data_to_request_parses_string_metadata():
 
 @pytest.mark.asyncio
 async def test_add_litellm_data_to_request_user_spend_and_budget():
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     request_mock = MagicMock(spec=Request)
     request_mock.url.path = "/v1/completions"
@@ -240,8 +243,7 @@ async def test_add_litellm_data_to_request_user_spend_and_budget():
 
 @pytest.mark.asyncio
 async def test_add_litellm_data_to_request_audio_transcription_multipart():
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     # Setup request mock for /v1/audio/transcriptions
     request_mock = MagicMock(spec=Request)
@@ -306,8 +308,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks():
     """
     Test that litellm_disabled_callbacks from key metadata is properly added to the request data.
     """
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     # Setup mock request
     request_mock = MagicMock(spec=Request)
@@ -360,8 +361,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks_empty():
     """
     Test that litellm_disabled_callbacks is not added when it's empty.
     """
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     # Setup mock request
     request_mock = MagicMock(spec=Request)
@@ -413,8 +413,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks_not_present():
     """
     Test that litellm_disabled_callbacks is not added when it's not present in metadata.
     """
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     # Setup mock request
     request_mock = MagicMock(spec=Request)
@@ -466,8 +465,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks_invalid_type():
     """
     Test that litellm_disabled_callbacks is not added when it's not a list.
     """
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     # Setup mock request
     request_mock = MagicMock(spec=Request)
@@ -519,8 +517,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks_with_logging_setti
     """
     Test that litellm_disabled_callbacks works correctly alongside logging settings.
     """
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     # Setup mock request
     request_mock = MagicMock(spec=Request)
@@ -1030,8 +1027,7 @@ from unittest.mock import AsyncMock
 from fastapi.responses import Response
 
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.proxy.common_request_processing import \
-    ProxyBaseLLMRequestProcessing
+from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 from litellm.proxy.utils import ProxyLogging
 from litellm.types.utils import StandardLoggingPayload
 
@@ -1149,6 +1145,47 @@ async def test_add_litellm_metadata_from_request_headers():
         litellm.callbacks = original_callbacks
 
 
+def test_add_litellm_metadata_from_request_headers_x_litellm_trace_id_sets_chain_id():
+    """x-litellm-trace-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining."""
+    headers = {"x-litellm-trace-id": "foo"}
+    data = {"metadata": {}}
+    LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+        headers=headers, data=data, _metadata_variable_name="metadata"
+    )
+    assert data["metadata"]["trace_id"] == "foo"
+    assert data["metadata"]["session_id"] == "foo"
+    assert data["litellm_session_id"] == "foo"
+    assert data["litellm_trace_id"] == "foo"
+
+
+def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id():
+    """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining."""
+    headers = {"x-litellm-session-id": "bar"}
+    data = {"metadata": {}}
+    LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+        headers=headers, data=data, _metadata_variable_name="metadata"
+    )
+    assert data["metadata"]["trace_id"] == "bar"
+    assert data["metadata"]["session_id"] == "bar"
+    assert data["litellm_session_id"] == "bar"
+    assert data["litellm_trace_id"] == "bar"
+
+
+def test_add_litellm_metadata_from_request_headers_both_headers_trace_id_precedence():
+    """When both x-litellm-trace-id and x-litellm-session-id are present, trace-id takes precedence for chain_id."""
+    headers = {
+        "x-litellm-trace-id": "trace-value",
+        "x-litellm-session-id": "session-value",
+    }
+    data = {"metadata": {}}
+    LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+        headers=headers, data=data, _metadata_variable_name="metadata"
+    )
+    assert data["metadata"]["trace_id"] == "trace-value"
+    assert data["metadata"]["session_id"] == "trace-value"
+    assert data["litellm_session_id"] == "trace-value"
+    assert data["litellm_trace_id"] == "trace-value"
+
 
 def test_get_internal_user_header_from_mapping_returns_expected_header():
     mappings = [
@@ -1407,8 +1444,7 @@ async def test_embedding_header_forwarding_with_model_group():
     importlib.reload(pre_call_utils_module)
 
     # Re-import the function after reload to get the fresh version
-    from litellm.proxy.litellm_pre_call_utils import \
-        add_litellm_data_to_request
+    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
 
     # Setup mock request for embeddings
     request_mock = MagicMock(spec=Request)
@@ -1542,11 +1578,13 @@ async def test_add_guardrails_from_policy_engine():
     Test that add_guardrails_from_policy_engine adds guardrails from matching policies
     and tracks applied policies in metadata.
     """
-    from litellm.proxy.policy_engine.attachment_registry import \
-        get_attachment_registry
+    from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
     from litellm.proxy.policy_engine.policy_registry import get_policy_registry
-    from litellm.types.proxy.policy_engine import (Policy, PolicyAttachment,
-                                                   PolicyGuardrails)
+    from litellm.types.proxy.policy_engine import (
+        Policy,
+        PolicyAttachment,
+        PolicyGuardrails,
+    )
 
     # Setup test data
     data = {
@@ -1659,8 +1697,7 @@ async def test_add_guardrails_from_policy_engine_policy_version_by_id():
     Test that add_guardrails_from_policy_engine executes a specific policy version
     when policy_ is passed in the request body.
     """
-    from litellm.proxy.policy_engine.attachment_registry import \
-        get_attachment_registry
+    from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
     from litellm.proxy.policy_engine.policy_registry import get_policy_registry
     from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails
 
@@ -1729,6 +1766,7 @@ async def test_bearer_token_not_in_debug_logs():
     """
     import logging
     from io import StringIO
+
     from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
     from litellm.proxy.proxy_server import ProxyConfig
 
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 b1012642c38..036a24c045a 100644
--- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx
@@ -6,9 +6,9 @@ import {
   LeftOutlined,
   RightOutlined,
 } from "@ant-design/icons";
-import { Sparkles, Wrench } from "lucide-react";
+import { Bot, Sparkles, Wrench } from "lucide-react";
 import { LogEntry } from "../columns";
-import { MCP_CALL_TYPES } from "../constants";
+import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "../constants";
 import { getEventDisplayName } from "../utils";
 import { DrawerHeader } from "./DrawerHeader";
 import { useKeyboardNavigation } from "./useKeyboardNavigation";
@@ -46,6 +46,7 @@ interface TraceEventRowProps {
 
 function TraceEventRow({ row, isSelected, onClick }: TraceEventRowProps) {
   const isMcp = MCP_CALL_TYPES.includes(row.call_type);
+  const isAgent = AGENT_CALL_TYPES.includes(row.call_type);
   const durationValue =
     row.request_duration_ms != null
       ? (row.request_duration_ms / 1000).toFixed(3)
@@ -64,6 +65,8 @@ function TraceEventRow({ row, isSelected, onClick }: TraceEventRowProps) {
       

         {isMcp ? (
           
+        ) : isAgent ? (
+          
         ) : (
           
         )}
@@ -219,7 +222,10 @@ export function LogDetailsDrawer({
     : null;
   const sessionDurationSeconds =
     sessionStart && sessionEnd ? ((sessionEnd.getTime() - sessionStart.getTime()) / 1000).toFixed(2) : "0.00";
-  const llmCount = sessionLogs.filter((row) => !MCP_CALL_TYPES.includes(row.call_type)).length;
+  const llmCount = sessionLogs.filter(
+    (row) => !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type),
+  ).length;
+  const agentCount = sessionLogs.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length;
   const mcpCount = sessionLogs.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length;
   const logsForList = isSessionMode ? sessionLogs : currentLog ? [currentLog] : [];
   const leftPanelId = isSessionMode ? sessionId || "" : currentLog?.request_id || "";
@@ -302,14 +308,25 @@ export function LogDetailsDrawer({
               

               

                 {logsForList.length} req
-                ·
-                {isSessionMode
-                  ? `${llmCount} LLM`
-                  : `${logsForList.filter((row) => !MCP_CALL_TYPES.includes(row.call_type)).length} LLM`}
-                ·
-                {isSessionMode
-                  ? `${mcpCount} MCP`
-                  : `${logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length} MCP`}
+                {[
+                  isSessionMode
+                    ? llmCount
+                    : logsForList.filter(
+                        (row) =>
+                          !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type),
+                      ).length,
+                  isSessionMode ? agentCount : logsForList.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length,
+                  isSessionMode ? mcpCount : logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length,
+                ].map((count, i) => {
+                  const label = [" LLM", " Agent", " MCP"][i];
+                  return count > 0 ? (
+                    
+                      ·
+                      {count}
+                      {label}
+                    
+                  ) : null;
+                })}
                 ·
                 {isSessionMode
                   ? getSpendString(totalSessionCost)
diff --git a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx
index e4195ece9ec..70c71c7f255 100644
--- a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx
@@ -1,5 +1,5 @@
 /**
- * Compact type-indicator badges for LLM and MCP log entries.
+ * Compact type-indicator badges for LLM, Agent, and MCP log entries.
  * Used in the request logs table and session type column.
  */
 
@@ -15,6 +15,18 @@ export const WrenchIcon = ({ size = 10 }: { size?: number }) => (
   
 );
 
+/** Agent/bot icon for A2A and agent call types (Lucide Bot-style). */
+export const AgentIcon = ({ size = 12 }: { size?: number }) => (
+  
+    
+    
+    
+    
+    
+    
+  
+);
+
 export const LlmBadge = ({ count }: { count?: number }) => (
   
     
@@ -28,3 +40,10 @@ export const McpBadge = ({ count }: { count?: number }) => (
     {count != null ? count : "MCP"}
   
 );
+
+export const AgentBadge = ({ count }: { count?: number }) => (
+  
+    
+    {count != null ? count : "Agent"}
+  
+);
diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx
index 7cea1a36383..3d9d73bb09b 100644
--- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx
@@ -6,8 +6,8 @@ import React, { useState } from "react";
 import { getProviderLogoAndName } from "../provider_info_helpers";
 import { TableHeaderSortDropdown } from "../common_components/TableHeaderSortDropdown/TableHeaderSortDropdown";
 import { TimeCell } from "./time_cell";
-import { MCP_CALL_TYPES } from "./constants";
-import { LlmBadge, McpBadge, SparkleIcon, WrenchIcon } from "./TypeBadges";
+import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "./constants";
+import { AgentBadge, AgentIcon, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } from "./TypeBadges";
 
 /** API sort field mapping for /spend/logs/ui endpoint */
 export const LOGS_SORT_FIELD_MAP = {
@@ -69,6 +69,7 @@ export type LogEntry = {
   mcp_tool_call_spend?: number;
   session_llm_count?: number;
   session_mcp_count?: number;
+  session_agent_count?: number;
   onKeyHashClick?: (keyHash: string) => void;
   onSessionClick?: (sessionId: string) => void;
 };
@@ -124,17 +125,26 @@ export const createColumns = (sortProps?: LogsSortProps): ColumnDef[]
       const row = info.row.original;
       const sessionCount = row.session_total_count || 1;
       const isMcp = MCP_CALL_TYPES.includes(row.call_type);
-      const sessionLlmCount = row.session_llm_count ?? (isMcp ? 0 : sessionCount);
+      const isAgent = AGENT_CALL_TYPES.includes(row.call_type);
+      const sessionLlmCount = row.session_llm_count ?? (isMcp || isAgent ? 0 : sessionCount);
+      const sessionAgentCount = row.session_agent_count ?? (isAgent ? sessionCount : 0);
       const sessionMcpCount = row.session_mcp_count ?? (isMcp ? sessionCount : 0);
 
       if (isMcp) return ;
+      if (isAgent && sessionCount <= 1) return ;
       if (sessionCount <= 1) return ;
 
-      // Multi-call session — show total count, plus MCP indicator when mixed.
+      // Multi-call session — show total count, plus Agent/MCP indicators when mixed.
       const sessionTypeBadge = (
         
           
           {sessionCount}
+          {sessionAgentCount > 0 && (
+            <>
+              ·
+              
+            
+          )}
           {sessionMcpCount > 0 && (
             <>
               ·
@@ -144,8 +154,13 @@ export const createColumns = (sortProps?: LogsSortProps): ColumnDef[]
         
       );
 
+      const tooltipParts = [
+        sessionLlmCount > 0 && `${sessionLlmCount} LLM`,
+        sessionAgentCount > 0 && `${sessionAgentCount} Agent`,
+        sessionMcpCount > 0 && `${sessionMcpCount} MCP`,
+      ].filter(Boolean);
       return (
-        
+        
           {sessionTypeBadge}
         
       );
diff --git a/ui/litellm-dashboard/src/components/view_logs/constants.ts b/ui/litellm-dashboard/src/components/view_logs/constants.ts
index 949dab275fe..57155feae23 100644
--- a/ui/litellm-dashboard/src/components/view_logs/constants.ts
+++ b/ui/litellm-dashboard/src/components/view_logs/constants.ts
@@ -15,6 +15,9 @@ export const ERROR_CODE_OPTIONS: { label: string; value: string }[] = [
 /** Call types that represent MCP tool invocations (shared across columns, index, drawer). */
 export const MCP_CALL_TYPES = ["call_mcp_tool", "list_mcp_tools"];
 
+/** Call types that represent agent/A2A requests (e.g. asend_message). */
+export const AGENT_CALL_TYPES = ["asend_message"];
+
 export const QUICK_SELECT_OPTIONS: { label: string; value: number; unit: string }[] = [
   { label: "Last 15 Minutes", value: 15, unit: "minutes" },
   { label: "Last Hour", value: 1, unit: "hours" },
diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx
index 64dc53c0e77..ab70126a5a8 100644
--- a/ui/litellm-dashboard/src/components/view_logs/index.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx
@@ -20,7 +20,7 @@ import KeyInfoView from "../templates/key_info_view";
 import AuditLogs from "./audit_logs";
 import { createColumns, LogEntry, type LogsSortField } from "./columns";
 import { ConfigInfoMessage } from "./ConfigInfoMessage";
-import { ERROR_CODE_OPTIONS, MCP_CALL_TYPES, QUICK_SELECT_OPTIONS } from "./constants";
+import { AGENT_CALL_TYPES, ERROR_CODE_OPTIONS, MCP_CALL_TYPES, QUICK_SELECT_OPTIONS } from "./constants";
 import { CostBreakdownViewer } from "./CostBreakdownViewer";
 import { ErrorViewer } from "./ErrorViewer";
 import { useLogFilterLogic } from "./log_filter_logic";
@@ -309,13 +309,15 @@ export default function SpendLogsTable({
     return matchesSearch;
   });
 
-  const sessionCompositionById = searchedLogs.reduce>((acc, log) => {
+  const sessionCompositionById = searchedLogs.reduce>((acc, log) => {
     if (!log.session_id) return acc;
     if (!acc[log.session_id]) {
-      acc[log.session_id] = { llm: 0, mcp: 0 };
+      acc[log.session_id] = { llm: 0, agent: 0, mcp: 0 };
     }
     if (MCP_CALL_TYPES.includes(log.call_type)) {
       acc[log.session_id].mcp += 1;
+    } else if (AGENT_CALL_TYPES.includes(log.call_type)) {
+      acc[log.session_id].agent += 1;
     } else {
       acc[log.session_id].llm += 1;
     }
@@ -343,6 +345,7 @@ export default function SpendLogsTable({
           request_duration_ms: log.request_duration_ms,
           session_llm_count: sessionComposition?.llm ?? undefined,
           session_mcp_count: sessionComposition?.mcp ?? undefined,
+          session_agent_count: sessionComposition?.agent ?? undefined,
           onKeyHashClick: (keyHash: string) => setSelectedKeyIdInfoView(keyHash),
           onSessionClick: (sessionId: string) => {
             if (sessionId) {
diff --git a/ui/litellm-dashboard/tsconfig.json b/ui/litellm-dashboard/tsconfig.json
index d24bdd340f7..5b0352feb98 100644
--- a/ui/litellm-dashboard/tsconfig.json
+++ b/ui/litellm-dashboard/tsconfig.json
@@ -14,7 +14,7 @@
     "moduleResolution": "bundler",
     "resolveJsonModule": true,
     "isolatedModules": true,
-    "jsx": "react-jsx",
+    "jsx": "preserve",
     "incremental": true,
     "plugins": [
       {

From 8baa3ae8cb5d03df43ee53e6317a4d5781ac2ac1 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff 
Date: Tue, 3 Mar 2026 19:56:55 -0800
Subject: [PATCH 125/147] [Feat] UI - Add Open in New Tab on leftnav Bar
 (#22731)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Add minimal dev_config.yaml for proxy development

Co-authored-by: Ishaan Jaff 

* feat(ui): wrap left nav items in  tags for open-in-new-tab support

Nav items are now rendered as  elements with proper href attributes,
enabling right-click → 'Open in new tab', Ctrl/Cmd+click, and
middle-click to open any sidebar page in a new browser tab.

Normal clicks continue to use SPA navigation (no full page reload).

Applied to both leftnav.tsx (query-param routing) and Sidebar2.tsx
(Next.js file-based routing).

Co-authored-by: Ishaan Jaff 

---------

Co-authored-by: Cursor Agent 
Co-authored-by: Ishaan Jaff 
---
 dev_config.yaml                               | 13 ++++++
 .../app/(dashboard)/components/Sidebar2.tsx   | 25 ++++++++++-
 .../src/components/leftnav.tsx                | 44 ++++++++++++++++++-
 3 files changed, 78 insertions(+), 4 deletions(-)
 create mode 100644 dev_config.yaml

diff --git a/dev_config.yaml b/dev_config.yaml
new file mode 100644
index 00000000000..64e3c14703e
--- /dev/null
+++ b/dev_config.yaml
@@ -0,0 +1,13 @@
+model_list:
+  - model_name: fake-openai-endpoint
+    litellm_params:
+      model: openai/fake-model
+      api_key: fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+
+general_settings:
+  master_key: sk-1234
+
+litellm_settings:
+  drop_params: True
+  telemetry: False
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx
index a74d3c108d6..b3829d0a8f4 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx
@@ -374,6 +374,27 @@ const Sidebar2: React.FC = ({ accessToken, userRole, defaultSelect
     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): React.ReactNode => {
+    const href = toHref(page);
+    return (
+       {
+          if (e.metaKey || e.ctrlKey || e.shiftKey || e.button === 1) {
+            e.stopPropagation();
+            return;
+          }
+          e.preventDefault();
+        }}
+        style={{ color: "inherit", textDecoration: "none" }}
+      >
+        {label}
+      
+    );
+  };
+
   return (
     
        = ({ accessToken, userRole, defaultSelect
             items={filteredMenuItems.map((item) => ({
               key: item.key,
               icon: item.icon,
-              label: item.label,
+              label: renderNavLink(item.label, item.page),
               children: item.children?.map((child) => ({
                 key: child.key,
                 icon: child.icon,
-                label: child.label,
+                label: renderNavLink(child.label, child.page),
                 onClick: () => goTo(child.page),
               })),
               onClick: !item.children ? () => goTo(item.page) : undefined,
diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx
index e84b6e86e4a..e09a95d15b7 100644
--- a/ui/litellm-dashboard/src/components/leftnav.tsx
+++ b/ui/litellm-dashboard/src/components/leftnav.tsx
@@ -374,6 +374,46 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse
     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.
+  const renderNavLink = (
+    label: React.ReactNode,
+    page: string,
+    externalUrl?: string,
+  ): React.ReactNode => {
+    if (externalUrl) {
+      return (
+         e.stopPropagation()}
+          style={{ color: "inherit", textDecoration: "none" }}
+        >
+          {label}
+        
+      );
+    }
+    const params = new URLSearchParams(window.location.search);
+    params.set("page", page);
+    const href = `?${params.toString()}`;
+    return (
+       {
+          if (e.metaKey || e.ctrlKey || e.shiftKey || e.button === 1) {
+            e.stopPropagation();
+            return;
+          }
+          e.preventDefault();
+        }}
+        style={{ color: "inherit", textDecoration: "none" }}
+      >
+        {label}
+      
+    );
+  };
+
   // Filter items based on user role and enabled pages for internal users
   const filterItemsByRole = (items: MenuItem[]): MenuItem[] => {
     const isAdmin = isAdminRole(userRole);
@@ -469,11 +509,11 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse
         children: filteredItems.map((item) => ({
           key: item.key,
           icon: item.icon,
-          label: item.label,
+          label: renderNavLink(item.label, item.page, item.external_url),
           children: item.children?.map((child) => ({
             key: child.key,
             icon: child.icon,
-            label: child.label,
+            label: renderNavLink(child.label, child.page, child.external_url),
             onClick: () => {
               if (child.external_url) {
                 window.open(child.external_url, "_blank");

From 1f412bc6d840e9b567ce07c8b26db3bf36008de7 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff 
Date: Tue, 3 Mar 2026 20:22:20 -0800
Subject: [PATCH 126/147] [Feat] Add Tool Policies for AI Gateway  (#22732)

* fix: fix ui render

* fix: fix minor bugs

* refactor: use prisma functions instead of raw sql (safer)

* fix(add-new-tiles-to-tool-policies): allow developer to see what's available

* feat: ensure tool allowlist runs correctly for tool names + mcp's

* refactor: more ui improvements

* feat: working key tool blocking

* feat(tools): show tool logs

* refactor: backend code improvements

* refactor: improve log viewer for tools

* fix: address PR review feedback for tool access control

- Add missing blocked_tools column to root schema.prisma (schema drift)
- Invalidate ToolPolicyRegistry after policy mutations so changes take effect immediately
- Remove dead code: unused get_effective_policies, get_tool_policies_cached, and helpers

Co-Authored-By: Claude Opus 4.6 

* fix: race condition in permission resolution and remove duplicate allowlist check

- Use atomic update_many with object_permission_id=None to prevent concurrent
  requests from creating orphaned permission rows and losing tool blocks
- Remove duplicate allowed_tools enforcement from guardrail (already enforced
  in auth layer via check_tools_allowlist)
- Move inline uuid import to module level

Co-Authored-By: Claude Opus 4.6 

* update to account for  userAgent

* UI - Add ToolDetails

* input/output policy

* LiteLLM_PolicyAttachmentTable

* LiteLLM_PolicyAttachmentTable

* fix: add _enqueue_tool_registry_upsert

* fix: tool mgmt endpoints

* tool mgmt endpoints

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update tests/test_litellm/proxy/db/test_tool_registry_writer.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* fix: sync root schema.prisma and fix test_tool_registry_writer for input/output policy

- Migrate root schema.prisma LiteLLM_ToolTable from call_policy to
  input_policy/output_policy, add missing user_agent and last_used_at columns
  (now consistent with litellm/proxy/schema.prisma and litellm-proxy-extras)
- Fix SpendLogToolIndex comment across all three schema files
- Fix all call_policy references in test_tool_registry_writer.py:
  swapped update_tool_policy arguments, wrong get_tools_by_names return type
  assertions, _mock_tool_row setting call_policy instead of input_policy

Addresses Greptile review feedback on PR #22732.

Made-with: Cursor

---------

Co-authored-by: Krrish Dholakia 
Co-authored-by: Claude Opus 4.6 
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 AGENTS.md                                     |   3 +
 CLAUDE.md                                     |   4 +
 .../migration.sql                             |   2 +
 .../migration.sql                             |  11 +
 .../litellm_proxy_extras/schema.prisma        |  46 +-
 .../chat/guardrail_translation/handler.py     |  10 +-
 .../guardrail_translation/base_translation.py |   7 +
 .../chat/guardrail_translation/handler.py     |  13 +
 .../guardrail_translation/handler.py          |  39 +-
 ...odel_prices_and_context_window_backup.json | 418 +++++++++++++--
 litellm/proxy/_new_secret_config.yaml         |  30 +-
 litellm/proxy/_types.py                       |   9 +-
 litellm/proxy/auth/auth_checks.py             |  55 +-
 litellm/proxy/db/db_spend_update_writer.py    | 155 +++---
 litellm/proxy/db/spend_log_tool_index.py      | 147 +++++
 litellm/proxy/db/tool_registry_writer.py      | 409 +++++++++++---
 .../tool_policy/tool_policy_guardrail.py      | 208 +++++---
 .../proxy/guardrails/tool_name_extraction.py  |  85 +++
 litellm/proxy/litellm_pre_call_utils.py       |   9 +
 .../tool_management_endpoints.py              | 501 +++++++++++++++++-
 litellm/proxy/proxy_server.py                 |  31 +-
 litellm/proxy/schema.prisma                   |  45 +-
 litellm/proxy/utils.py                        |  18 +-
 litellm/types/tool_management.py              |  66 ++-
 schema.prisma                                 |  45 +-
 scripts/test_tool_allowlist_script.py         | 116 ++++
 .../proxy/db/test_tool_registry_writer.py     | 272 ++++++----
 .../test_tool_policy_guardrail.py             |  85 +--
 .../proxy/test_tools_allowlist_enforcement.py | 200 +++++++
 ui/litellm-dashboard/src/app/page.tsx         |   4 +-
 .../src/components/ToolDetail.tsx             | 445 ++++++++++++++++
 .../src/components/ToolPolicies.tsx           | 304 +++++++----
 .../components/ToolPolicies/PolicySelect.tsx  |  92 ++++
 .../src/components/ToolPoliciesView.tsx       |  44 ++
 .../src/components/networking.tsx             | 166 +++++-
 35 files changed, 3467 insertions(+), 627 deletions(-)
 create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql
 create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql
 create mode 100644 litellm/proxy/db/spend_log_tool_index.py
 create mode 100644 litellm/proxy/guardrails/tool_name_extraction.py
 create mode 100644 scripts/test_tool_allowlist_script.py
 create mode 100644 tests/test_litellm/proxy/test_tools_allowlist_enforcement.py
 create mode 100644 ui/litellm-dashboard/src/components/ToolDetail.tsx
 create mode 100644 ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.tsx
 create mode 100644 ui/litellm-dashboard/src/components/ToolPoliciesView.tsx

diff --git a/AGENTS.md b/AGENTS.md
index d43f41dbe30..546f2997bf5 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -109,6 +109,8 @@ Key files:
 - `litellm/proxy/auth/` - Authentication logic
 - `litellm/proxy/management_endpoints/` - Admin API endpoints
 
+**Database (proxy)**: Use Prisma model methods (`prisma_client.db..upsert`, `.find_many`, `.find_unique`, etc.), not raw SQL (`execute_raw`/`query_raw`). See COMMON PITFALLS for details.
+
 ## MCP (MODEL CONTEXT PROTOCOL) SUPPORT
 
 LiteLLM supports MCP for agent workflows:
@@ -176,6 +178,7 @@ When opening issues or pull requests, follow these templates:
 5. **Dependencies**: Keep dependencies minimal and well-justified
 6. **UI/Backend Contract Mismatch**: When adding a new entity type to the UI, always check whether the backend endpoint accepts a single value or an array. Match the UI control accordingly (single-select vs. multi-select) to avoid silently dropping user selections
 7. **Missing Tests for New Entity Types**: When adding a new entity type (e.g., in `EntityUsage`, `UsageViewSelect`), always add corresponding tests in the existing test files and update any icon/component mocks
+8. **Raw SQL in proxy DB code**: Do not use `execute_raw` or `query_raw` for proxy database access. Use Prisma model methods (e.g. `prisma_client.db.litellm_tooltable.upsert()`, `.find_many()`, `.find_unique()`) so behavior stays consistent with the schema, the client stays mockable in tests, and you avoid the pitfalls of hand-written SQL (parameter ordering, type casting, schema drift)
 
 8. **Do not hardcode model-specific flags**: Put model-specific capability flags in `model_prices_and_context_window.json` and read them via `get_model_info` (or existing helpers like `supports_reasoning`). This prevents users from needing to upgrade LiteLLM each time a new model supports a feature.
 
diff --git a/CLAUDE.md b/CLAUDE.md
index 3b597fb8a90..c1eb75d2515 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -107,6 +107,10 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components:
 - Migration files auto-generated with `prisma migrate dev`
 - Always test migrations against both PostgreSQL and SQLite
 
+### Proxy database access
+- **Do not write raw SQL** for proxy DB operations. Use Prisma model methods instead of `execute_raw` / `query_raw`.
+- Use the generated client: `prisma_client.db.` (e.g. `litellm_tooltable`, `litellm_usertable`) with `.upsert()`, `.find_many()`, `.find_unique()`, `.update()`, `.update_many()` as appropriate. This avoids schema/client drift, keeps code testable with simple mocks, and matches patterns used in spend logs and other proxy code.
+
 ### Enterprise Features
 - Enterprise-specific code in `enterprise/` directory
 - Optional features enabled via environment variables
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql
new file mode 100644
index 00000000000..cba06684193
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "blocked_tools" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql
new file mode 100644
index 00000000000..e3199679ce2
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql
@@ -0,0 +1,11 @@
+-- CreateTable
+CREATE TABLE "LiteLLM_SpendLogToolIndex" (
+    "request_id" TEXT NOT NULL,
+    "tool_name" TEXT NOT NULL,
+    "start_time" TIMESTAMP(3) NOT NULL,
+
+    CONSTRAINT "LiteLLM_SpendLogToolIndex_pkey" PRIMARY KEY ("request_id","tool_name")
+);
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx" ON "LiteLLM_SpendLogToolIndex"("tool_name", "start_time");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
index e0b28a4e012..5abe7a0a2b1 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
@@ -260,6 +260,7 @@ model LiteLLM_ObjectPermissionTable {
   vector_stores         String[]       @default([])
   agents                String[]       @default([])
   agent_access_groups   String[]       @default([])
+  blocked_tools         String[]       @default([]) // Tool names blocked for any key/team/user with this permission
   teams                 LiteLLM_TeamTable[]
   projects              LiteLLM_ProjectTable[]
   verification_tokens   LiteLLM_VerificationToken[]
@@ -928,6 +929,16 @@ model LiteLLM_SpendLogGuardrailIndex {
   @@index([policy_id, start_time])
 }
 
+// Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production
+model LiteLLM_SpendLogToolIndex {
+  request_id String
+  tool_name  String   // matches LiteLLM_ToolTable.tool_name; join for input_policy/output_policy etc.
+  start_time DateTime
+
+  @@id([request_id, tool_name])
+  @@index([tool_name, start_time])
+}
+
 // Prompt table for storing prompt configurations
 model LiteLLM_PromptTable {
   id String @id @default(uuid())
@@ -1065,26 +1076,31 @@ model LiteLLM_PolicyAttachmentTable {
   updated_by       String?
 }
 
-// Global tool registry - auto-discovered from LLM responses; admins set call_policy here
+// Global tool registry - auto-discovered from LLM responses; admins set input_policy/output_policy here
 model LiteLLM_ToolTable {
-  tool_id     String   @id @default(uuid())
-  tool_name   String   @unique   // e.g. "huggingface_remote-mcp__dynamic_space"
-  origin      String?            // MCP server name or "user_defined"
-  call_policy String   @default("untrusted")  // "trusted" | "untrusted" | "dual_llm" | "blocked"
-  call_count  Int      @default(0)            // cumulative number of times this tool was seen
-  assignments Json?    @default("{}")
-  key_hash    String?            // hash of the virtual key that first called this tool
-  team_id     String?            // team that first called this tool
-  key_alias   String?            // human-readable alias of the virtual key
-  created_at  DateTime @default(now())
-  created_by  String?
-  updated_at  DateTime @default(now()) @updatedAt
-  updated_by  String?
+  tool_id       String    @id @default(uuid())
+  tool_name     String    @unique   // e.g. "huggingface_remote-mcp__dynamic_space"
+  origin        String?             // MCP server name or "user_defined"
+  input_policy  String    @default("untrusted")  // "trusted" | "untrusted" | "blocked"
+  output_policy String    @default("untrusted")  // "trusted" | "untrusted"
+  call_count    Int       @default(0)             // cumulative number of times this tool was seen
+  assignments   Json?     @default("{}")
+  key_hash      String?             // hash of the virtual key that first called this tool
+  team_id       String?             // team that first called this tool
+  key_alias     String?             // human-readable alias of the virtual key
+  user_agent    String?             // user-agent of the first request that discovered this tool
+  last_used_at  DateTime?           // timestamp of the most recent call
+  created_at    DateTime  @default(now())
+  created_by    String?
+  updated_at    DateTime  @default(now()) @updatedAt
+  updated_by    String?
 
-  @@index([call_policy])
+  @@index([input_policy])
+  @@index([output_policy])
   @@index([team_id])
 }
 
+// Per-(tool, team/key) policy overrides. When present, override replaces global tool policy for that scope.
 //Unified Access Groups table for storing unified access groups
 model LiteLLM_AccessGroupTable {
   access_group_id    String   @id @default(uuid())
diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py
index 98650a238e9..a6df346e8a8 100644
--- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py
+++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py
@@ -75,7 +75,7 @@ class AnthropicMessagesHandler(BaseTranslation):
         if messages is None:
             return data
 
-        chat_completion_compatible_request, tool_name_mapping = (
+        chat_completion_compatible_request, _tool_name_mapping = (
             LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
                 # Use a shallow copy to avoid mutating request data (pop on litellm_metadata).
                 anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
@@ -141,6 +141,14 @@ class AnthropicMessagesHandler(BaseTranslation):
 
         return data
 
+    def extract_request_tool_names(self, data: dict) -> List[str]:
+        """Extract tool names from Anthropic messages request (tools[].name)."""
+        names: List[str] = []
+        for tool in data.get("tools") or []:
+            if isinstance(tool, dict) and tool.get("name"):
+                names.append(str(tool["name"]))
+        return names
+
     def _extract_input_text_and_images(
         self,
         message: Dict[str, Any],
diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py
index 7106c207bd6..a7982cb606e 100644
--- a/litellm/llms/base_llm/guardrail_translation/base_translation.py
+++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py
@@ -98,3 +98,10 @@ class BaseTranslation(ABC):
         Optional to override in subclasses.
         """
         return responses_so_far
+
+    def extract_request_tool_names(self, data: dict) -> List[str]:
+        """
+        Extract tool names from the request body for allowlist/policy checks.
+        Override in tool-capable handlers; default returns [].
+        """
+        return []
diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py
index 67e9e42bc30..10b0b58b6ac 100644
--- a/litellm/llms/openai/chat/guardrail_translation/handler.py
+++ b/litellm/llms/openai/chat/guardrail_translation/handler.py
@@ -135,6 +135,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
 
         return data
 
+    def extract_request_tool_names(self, data: dict) -> List[str]:
+        """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
+        names: List[str] = []
+        for tool in data.get("tools") or []:
+            if isinstance(tool, dict) and tool.get("type") == "function":
+                fn = tool.get("function")
+                if isinstance(fn, dict) and fn.get("name"):
+                    names.append(str(fn["name"]))
+        for fn in data.get("functions") or []:
+            if isinstance(fn, dict) and fn.get("name"):
+                names.append(str(fn["name"]))
+        return names
+
     def _extract_inputs(
         self,
         message: Dict[str, Any],
diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py
index 6b092911d3c..7c3354cf88e 100644
--- a/litellm/llms/openai/responses/guardrail_translation/handler.py
+++ b/litellm/llms/openai/responses/guardrail_translation/handler.py
@@ -30,27 +30,22 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
 
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
 
-from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
+from openai.types.responses.response_function_tool_call import \
+    ResponseFunctionToolCall
 from pydantic import BaseModel
 
 from litellm._logging import verbose_proxy_logger
 from litellm.completion_extras.litellm_responses_transformation.transformation import (
     LiteLLMResponsesTransformationHandler,
-    OpenAiResponsesToChatCompletionStreamIterator,
-)
-from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
-from litellm.responses.litellm_completion_transformation.transformation import (
-    LiteLLMCompletionResponsesConfig,
-)
-from litellm.types.llms.openai import (
-    ChatCompletionToolCallChunk,
-    ChatCompletionToolParam,
-)
-from litellm.types.responses.main import (
-    GenericResponseOutputItem,
-    OutputFunctionToolCall,
-    OutputText,
-)
+    OpenAiResponsesToChatCompletionStreamIterator)
+from litellm.llms.base_llm.guardrail_translation.base_translation import \
+    BaseTranslation
+from litellm.responses.litellm_completion_transformation.transformation import \
+    LiteLLMCompletionResponsesConfig
+from litellm.types.llms.openai import (ChatCompletionToolCallChunk,
+                                       ChatCompletionToolParam)
+from litellm.types.responses.main import (GenericResponseOutputItem,
+                                          OutputFunctionToolCall, OutputText)
 from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
@@ -188,6 +183,18 @@ class OpenAIResponsesHandler(BaseTranslation):
 
         return data
 
+    def extract_request_tool_names(self, data: dict) -> List[str]:
+        """Extract tool names from Responses API request (tools[].name for function, tools[].server_label for mcp)."""
+        names: List[str] = []
+        for tool in data.get("tools") or []:
+            if not isinstance(tool, dict):
+                continue
+            if tool.get("type") == "function" and tool.get("name"):
+                names.append(str(tool["name"]))
+            elif tool.get("type") == "mcp" and tool.get("server_label"):
+                names.append(str(tool["server_label"]))
+        return names
+
     def _extract_and_transform_tools(
         self,
         tools: List[Dict[str, Any]],
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 3dfbc5cb219..b92e2727979 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -9779,6 +9779,122 @@
             }
         ]
     },
+    "dashscope/qwen3-max-2026-01-23": {
+        "litellm_provider": "dashscope",
+        "max_input_tokens": 258048,
+        "max_output_tokens": 65536,
+        "max_tokens": 65536,
+        "mode": "chat",
+        "source": "https://www.alibabacloud.com/help/en/model-studio/models",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "tiered_pricing": [
+            {
+                "input_cost_per_token": 1.2e-06,
+                "output_cost_per_token": 6e-06,
+                "range": [
+                    0,
+                    32000.0
+                ]
+            },
+            {
+                "input_cost_per_token": 2.4e-06,
+                "output_cost_per_token": 1.2e-05,
+                "range": [
+                    32000.0,
+                    128000.0
+                ]
+            },
+            {
+                "input_cost_per_token": 3e-06,
+                "output_cost_per_token": 1.5e-05,
+                "range": [
+                    128000.0,
+                    252000.0
+                ]
+            }
+        ]
+    },
+    "dashscope/qwen3-next-80b-a3b-instruct": {
+        "input_cost_per_token": 1.5e-07,
+        "litellm_provider": "dashscope",
+        "max_input_tokens": 262144,
+        "max_output_tokens": 65536,
+        "max_tokens": 65536,
+        "mode": "chat",
+        "output_cost_per_token": 1.2e-06,
+        "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
+        "supports_function_calling": true,
+        "supports_tool_choice": true
+    },
+    "dashscope/qwen3-next-80b-a3b-thinking": {
+        "input_cost_per_token": 1.5e-07,
+        "litellm_provider": "dashscope",
+        "max_input_tokens": 262144,
+        "max_output_tokens": 65536,
+        "max_tokens": 65536,
+        "mode": "chat",
+        "output_cost_per_token": 1.2e-06,
+        "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true
+    },
+    "dashscope/qwen3-vl-235b-a22b-instruct": {
+        "input_cost_per_token": 4e-07,
+        "litellm_provider": "dashscope",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "output_cost_per_token": 1.6e-06,
+        "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
+        "supports_function_calling": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "dashscope/qwen3-vl-235b-a22b-thinking": {
+        "input_cost_per_token": 4e-07,
+        "litellm_provider": "dashscope",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "output_cost_per_token": 4e-06,
+        "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "dashscope/qwen3-vl-32b-instruct": {
+        "input_cost_per_token": 1.6e-07,
+        "litellm_provider": "dashscope",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "output_cost_per_token": 6.4e-07,
+        "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
+        "supports_function_calling": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "dashscope/qwen3-vl-32b-thinking": {
+        "input_cost_per_token": 1.6e-07,
+        "litellm_provider": "dashscope",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "output_cost_per_token": 2.87e-06,
+        "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "dashscope/qwen3-vl-plus": {
         "litellm_provider": "dashscope",
         "max_input_tokens": 260096,
@@ -10844,7 +10960,8 @@
         "output_cost_per_token": 9e-08,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": {
         "max_tokens": 131072,
@@ -10854,7 +10971,8 @@
         "output_cost_per_token": 1e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": {
         "max_tokens": 131072,
@@ -10874,7 +10992,8 @@
         "output_cost_per_token": 4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen2.5-72B-Instruct": {
         "max_tokens": 32768,
@@ -10884,7 +11003,8 @@
         "output_cost_per_token": 3.9e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen2.5-7B-Instruct": {
         "max_tokens": 32768,
@@ -10905,7 +11025,8 @@
         "litellm_provider": "deepinfra",
         "mode": "chat",
         "supports_tool_choice": true,
-        "supports_vision": true
+        "supports_vision": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-14B": {
         "max_tokens": 40960,
@@ -10915,7 +11036,8 @@
         "output_cost_per_token": 2.4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-235B-A22B": {
         "max_tokens": 40960,
@@ -10925,7 +11047,8 @@
         "output_cost_per_token": 5.4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": {
         "max_tokens": 262144,
@@ -10935,7 +11058,8 @@
         "output_cost_per_token": 6e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": {
         "max_tokens": 262144,
@@ -10945,7 +11069,8 @@
         "output_cost_per_token": 2.9e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-30B-A3B": {
         "max_tokens": 40960,
@@ -10955,7 +11080,8 @@
         "output_cost_per_token": 2.9e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-32B": {
         "max_tokens": 40960,
@@ -10965,7 +11091,8 @@
         "output_cost_per_token": 2.8e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": {
         "max_tokens": 262144,
@@ -10975,7 +11102,8 @@
         "output_cost_per_token": 1.6e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": {
         "max_tokens": 262144,
@@ -10985,7 +11113,8 @@
         "output_cost_per_token": 1.2e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-Next-80B-A3B-Instruct": {
         "max_tokens": 262144,
@@ -10995,7 +11124,8 @@
         "output_cost_per_token": 1.4e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Qwen/Qwen3-Next-80B-A3B-Thinking": {
         "max_tokens": 262144,
@@ -11005,7 +11135,8 @@
         "output_cost_per_token": 1.4e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": {
         "max_tokens": 8192,
@@ -11056,7 +11187,8 @@
         "cache_read_input_token_cost": 3.3e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/anthropic/claude-4-opus": {
         "max_tokens": 200000,
@@ -11066,7 +11198,8 @@
         "output_cost_per_token": 8.25e-05,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/anthropic/claude-4-sonnet": {
         "max_tokens": 200000,
@@ -11076,7 +11209,8 @@
         "output_cost_per_token": 1.65e-05,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-R1": {
         "max_tokens": 163840,
@@ -11086,7 +11220,8 @@
         "output_cost_per_token": 2.4e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-R1-0528": {
         "max_tokens": 163840,
@@ -11097,7 +11232,8 @@
         "cache_read_input_token_cost": 4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": {
         "max_tokens": 32768,
@@ -11107,7 +11243,8 @@
         "output_cost_per_token": 3e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": {
         "max_tokens": 131072,
@@ -11127,7 +11264,8 @@
         "output_cost_per_token": 2.7e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": {
         "max_tokens": 40960,
@@ -11137,7 +11275,8 @@
         "output_cost_per_token": 3e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-V3": {
         "max_tokens": 163840,
@@ -11147,7 +11286,8 @@
         "output_cost_per_token": 8.9e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-V3-0324": {
         "max_tokens": 163840,
@@ -11157,7 +11297,8 @@
         "output_cost_per_token": 8.8e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-V3.1": {
         "max_tokens": 163840,
@@ -11169,7 +11310,8 @@
         "litellm_provider": "deepinfra",
         "mode": "chat",
         "supports_tool_choice": true,
-        "supports_reasoning": true
+        "supports_reasoning": true,
+        "supports_function_calling": true
     },
     "deepinfra/deepseek-ai/DeepSeek-V3.1-Terminus": {
         "max_tokens": 163840,
@@ -11180,7 +11322,8 @@
         "cache_read_input_token_cost": 2.16e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/google/gemini-2.0-flash-001": {
         "deprecation_date": "2026-06-01",
@@ -11191,7 +11334,8 @@
         "output_cost_per_token": 4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/google/gemini-2.5-flash": {
         "max_tokens": 1000000,
@@ -11201,7 +11345,8 @@
         "output_cost_per_token": 2.5e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/google/gemini-2.5-pro": {
         "max_tokens": 1000000,
@@ -11211,7 +11356,8 @@
         "output_cost_per_token": 1e-05,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/google/gemma-3-12b-it": {
         "max_tokens": 131072,
@@ -11221,7 +11367,8 @@
         "output_cost_per_token": 1e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/google/gemma-3-27b-it": {
         "max_tokens": 131072,
@@ -11231,7 +11378,8 @@
         "output_cost_per_token": 1.6e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/google/gemma-3-4b-it": {
         "max_tokens": 131072,
@@ -11241,7 +11389,8 @@
         "output_cost_per_token": 8e-08,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": {
         "max_tokens": 131072,
@@ -11261,7 +11410,8 @@
         "output_cost_per_token": 2e-08,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Llama-3.3-70B-Instruct": {
         "max_tokens": 131072,
@@ -11271,7 +11421,8 @@
         "output_cost_per_token": 4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": {
         "max_tokens": 131072,
@@ -11281,6 +11432,7 @@
         "output_cost_per_token": 3.9e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
+        "supports_function_calling": true,
         "supports_tool_choice": true
     },
     "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": {
@@ -11291,7 +11443,8 @@
         "output_cost_per_token": 6e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": {
         "max_tokens": 327680,
@@ -11301,7 +11454,8 @@
         "output_cost_per_token": 3e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Llama-Guard-3-8B": {
         "max_tokens": 131072,
@@ -11331,7 +11485,8 @@
         "output_cost_per_token": 6e-08,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": {
         "max_tokens": 131072,
@@ -11341,7 +11496,8 @@
         "output_cost_per_token": 4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": {
         "max_tokens": 131072,
@@ -11351,7 +11507,8 @@
         "output_cost_per_token": 2.8e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": {
         "max_tokens": 131072,
@@ -11361,7 +11518,8 @@
         "output_cost_per_token": 5e-08,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": {
         "max_tokens": 131072,
@@ -11371,7 +11529,8 @@
         "output_cost_per_token": 3e-08,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/microsoft/WizardLM-2-8x22B": {
         "max_tokens": 65536,
@@ -11391,7 +11550,8 @@
         "output_cost_per_token": 1.4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": {
         "max_tokens": 131072,
@@ -11401,7 +11561,8 @@
         "output_cost_per_token": 4e-08,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": {
         "max_tokens": 32768,
@@ -11411,7 +11572,8 @@
         "output_cost_per_token": 8e-08,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": {
         "max_tokens": 128000,
@@ -11421,7 +11583,8 @@
         "output_cost_per_token": 2e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": {
         "max_tokens": 32768,
@@ -11431,7 +11594,8 @@
         "output_cost_per_token": 4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/moonshotai/Kimi-K2-Instruct": {
         "max_tokens": 131072,
@@ -11441,7 +11605,8 @@
         "output_cost_per_token": 2e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/moonshotai/Kimi-K2-Instruct-0905": {
         "max_tokens": 262144,
@@ -11452,7 +11617,8 @@
         "cache_read_input_token_cost": 4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": {
         "max_tokens": 131072,
@@ -11462,7 +11628,8 @@
         "output_cost_per_token": 6e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/nvidia/Llama-3.3-Nemotron-Super-49B-v1.5": {
         "max_tokens": 131072,
@@ -11472,7 +11639,8 @@
         "output_cost_per_token": 4e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/nvidia/NVIDIA-Nemotron-Nano-9B-v2": {
         "max_tokens": 131072,
@@ -11482,7 +11650,8 @@
         "output_cost_per_token": 1.6e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/openai/gpt-oss-120b": {
         "max_tokens": 131072,
@@ -11492,7 +11661,8 @@
         "output_cost_per_token": 4.5e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/openai/gpt-oss-20b": {
         "max_tokens": 131072,
@@ -11502,7 +11672,8 @@
         "output_cost_per_token": 1.5e-07,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepinfra/zai-org/GLM-4.5": {
         "max_tokens": 131072,
@@ -11512,7 +11683,8 @@
         "output_cost_per_token": 1.6e-06,
         "litellm_provider": "deepinfra",
         "mode": "chat",
-        "supports_tool_choice": true
+        "supports_tool_choice": true,
+        "supports_function_calling": true
     },
     "deepseek/deepseek-chat": {
         "cache_creation_input_token_cost": 0.0,
@@ -25806,6 +25978,30 @@
         "supports_vision": true,
         "tool_use_system_prompt_tokens": 159
     },
+    "openrouter/anthropic/claude-sonnet-4.6": {
+        "cache_creation_input_token_cost": 3.75e-06,
+        "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
+        "cache_read_input_token_cost": 3e-07,
+        "cache_read_input_token_cost_above_200k_tokens": 6e-07,
+        "input_cost_per_token": 3e-06,
+        "input_cost_per_token_above_200k_tokens": 6e-06,
+        "litellm_provider": "openrouter",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 1.5e-05,
+        "output_cost_per_token_above_200k_tokens": 2.25e-05,
+        "source": "https://openrouter.ai/anthropic/claude-sonnet-4.6",
+        "supports_assistant_prefill": true,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "tool_use_system_prompt_tokens": 159
+    },
     "openrouter/anthropic/claude-opus-4.5": {
         "cache_creation_input_token_cost": 6.25e-06,
         "cache_read_input_token_cost": 5e-07,
@@ -26156,6 +26352,39 @@
         "supports_web_search": true,
         "tpm": 800000
     },
+    "openrouter/google/gemini-3.1-pro-preview": {
+        "cache_read_input_token_cost": 2e-07,
+        "cache_read_input_token_cost_above_200k_tokens": 4e-07,
+        "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07,
+        "input_cost_per_token": 2e-06,
+        "input_cost_per_token_above_200k_tokens": 4e-06,
+        "litellm_provider": "openrouter",
+        "max_input_tokens": 1048576,
+        "max_output_tokens": 65536,
+        "max_tokens": 65536,
+        "mode": "chat",
+        "output_cost_per_token": 1.2e-05,
+        "output_cost_per_token_above_200k_tokens": 1.8e-05,
+        "source": "https://openrouter.ai/google/gemini-3.1-pro-preview",
+        "supported_modalities": [
+            "text",
+            "image",
+            "audio",
+            "video"
+        ],
+        "supported_output_modalities": [
+            "text"
+        ],
+        "supports_audio_input": true,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_system_messages": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "openrouter/gryphe/mythomax-l2-13b": {
         "input_cost_per_token": 1.875e-06,
         "litellm_provider": "openrouter",
@@ -26533,6 +26762,29 @@
         "supports_reasoning": true,
         "supports_tool_choice": true
     },
+    "openrouter/openai/gpt-5.1-codex-max": {
+        "cache_read_input_token_cost": 1.25e-07,
+        "input_cost_per_token": 1.25e-06,
+        "litellm_provider": "openrouter",
+        "max_input_tokens": 400000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 1e-05,
+        "source": "https://openrouter.ai/openai/gpt-5.1-codex-max",
+        "supported_modalities": [
+            "text",
+            "image"
+        ],
+        "supported_output_modalities": [
+            "text"
+        ],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "openrouter/openai/gpt-5.2": {
         "input_cost_per_image": 0,
         "cache_read_input_token_cost": 1.75e-07,
@@ -26687,6 +26939,19 @@
         "supports_tool_choice": true,
         "supports_function_calling": true
     },
+    "openrouter/qwen/qwen3-coder-plus": {
+        "input_cost_per_token": 1e-06,
+        "litellm_provider": "openrouter",
+        "max_input_tokens": 997952,
+        "max_output_tokens": 65536,
+        "max_tokens": 65536,
+        "mode": "chat",
+        "output_cost_per_token": 5e-06,
+        "source": "https://openrouter.ai/qwen/qwen3-coder-plus",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true
+    },
     "openrouter/qwen/qwen3-235b-a22b-2507": {
         "input_cost_per_token": 7.1e-08,
         "litellm_provider": "openrouter",
@@ -26822,6 +27087,19 @@
         "supports_vision": true,
         "supports_prompt_caching": false
     },
+    "openrouter/z-ai/glm-5": {
+        "input_cost_per_token": 8e-07,
+        "litellm_provider": "openrouter",
+        "max_input_tokens": 202752,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 2.56e-06,
+        "source": "https://openrouter.ai/z-ai/glm-5",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true
+    },
     "openrouter/minimax/minimax-m2.1": {
         "input_cost_per_token": 2.7e-07,
         "output_cost_per_token": 1.2e-06,
@@ -34327,6 +34605,36 @@
         "supports_tool_choice": true,
         "source": "https://aws.amazon.com/bedrock/pricing/"
     },
+    "zai/glm-5": {
+        "cache_creation_input_token_cost": 0,
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 1e-06,
+        "output_cost_per_token": 3.2e-06,
+        "litellm_provider": "zai",
+        "max_input_tokens": 200000,
+        "max_output_tokens": 128000,
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "source": "https://docs.z.ai/guides/overview/pricing"
+    },
+    "zai/glm-5-code": {
+        "cache_creation_input_token_cost": 0,
+        "cache_read_input_token_cost": 3e-07,
+        "input_cost_per_token": 1.2e-06,
+        "output_cost_per_token": 5e-06,
+        "litellm_provider": "zai",
+        "max_input_tokens": 200000,
+        "max_output_tokens": 128000,
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "source": "https://docs.z.ai/guides/overview/pricing"
+    },
     "zai/glm-4.7": {
         "cache_creation_input_token_cost": 0,
         "cache_read_input_token_cost": 1.1e-07,
diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml
index 6b84d90a327..508c1c94659 100644
--- a/litellm/proxy/_new_secret_config.yaml
+++ b/litellm/proxy/_new_secret_config.yaml
@@ -23,33 +23,11 @@ model_list:
 
 
 guardrails:
-  - guardrail_name: "airline-competitor-intent"
-    guardrail_id: "airline-competitor-intent"
+  - guardrail_name: "tool_policy"
     litellm_params:
-      guardrail: litellm_content_filter
-      mode: pre_call
-      default_on: false
-      competitor_intent_config:
-        brand_self:
-          - emirates
-          - ek
-        competitors:
-          - qatar airways
-          - qatar
-          - etihad
-        locations:
-          - qatar
-          - doha
-          - doh
-        competitor_aliases:
-          qatar airways: [qr, doha airline]
-          qatar: [qr]
-        policy:
-          competitor_comparison: refuse
-          possible_competitor_comparison: reframe
-        threshold_high: 0.70
-        threshold_medium: 0.45
-        threshold_low: 0.30
+      guardrail: tool_policy
+      mode: [pre_call, post_call]
+      default_on: true
 
 mcp_servers:
   my_http_server:
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 8d49020461d..42b48446e7a 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -77,6 +77,7 @@ class SupportedDBObjectType(str, enum.Enum):
     PASS_THROUGH_ENDPOINTS = "pass_through_endpoints"
     PROMPTS = "prompts"
     MODEL_COST_MAP = "model_cost_map"
+    TOOLS = "tools"
 
     def __str__(self):
         return str(self.value)
@@ -2133,7 +2134,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
     user_header_mappings: Optional[List[UserHeaderMapping]] = None
     supported_db_objects: Optional[List[SupportedDBObjectType]] = Field(
         None,
-        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'. If not set, all objects are loaded (default behavior).",
+        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'. If not set, all objects are loaded (default behavior).",
     )
     user_mcp_management_mode: Optional[UserMCPManagementMode] = Field(
         None,
@@ -3377,6 +3378,11 @@ class ProxyErrorTypes(str, enum.Enum):
     Team member is already in team
     """
 
+    tool_access_denied = "tool_access_denied"
+    """
+    Tool is not in the allowed tools list for this key/team
+    """
+
     @classmethod
     def get_model_access_error_type_for_object(
         cls, object_type: Literal["key", "user", "team", "org", "project"]
@@ -4161,6 +4167,7 @@ class ToolDiscoveryQueueItem(TypedDict, total=False):
     key_hash: Optional[str]  # hash of virtual key that triggered discovery
     team_id: Optional[str]  # team that triggered discovery
     key_alias: Optional[str]  # human-readable key alias
+    user_agent: Optional[str]  # HTTP User-Agent of the caller
 
 
 class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase):
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 91ac58215ab..79e5f78f68e 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -58,6 +58,10 @@ from litellm.proxy._types import (
 )
 from litellm.proxy.auth.route_checks import RouteChecks
 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.route_llm_request import route_request
 from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
 from litellm.router import Router
@@ -220,7 +224,48 @@ async def _run_project_checks(
         )
 
 
-async def common_checks(
+async def check_tools_allowlist(
+    request_body: dict,
+    valid_token: Optional[UserAPIKeyAuth],
+    team_object: Optional[LiteLLM_TeamTable],
+    route: str,
+) -> None:
+    """
+    Enforce key/team tool allowlist (metadata.allowed_tools). No DB in hot path —
+    effective allowlist is read from valid_token.metadata and valid_token.team_metadata.
+    Raises ProxyException with tool_access_denied if a tool is not allowed.
+    """
+    from litellm.litellm_core_utils.api_route_to_call_types import (
+        get_call_types_for_route,
+    )
+
+    if valid_token is None:
+        return
+    call_types = get_call_types_for_route(route)
+    if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types):
+        return
+    tool_names = extract_request_tool_names(route, request_body)
+    if not tool_names:
+        return
+    key_meta = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {}
+    team_meta = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {}
+    key_allowed = key_meta.get("allowed_tools")
+    team_allowed = team_meta.get("allowed_tools")
+    effective = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed
+    if not isinstance(effective, list) or len(effective) == 0:
+        return
+    allowed_set = {str(t) for t in effective}
+    disallowed = [n for n in tool_names if n not in allowed_set]
+    if disallowed:
+        raise ProxyException(
+            message=f"Tool(s) {disallowed} are not in the allowed tools list for this key/team.",
+            type=ProxyErrorTypes.tool_access_denied,
+            param="tools",
+            code=status.HTTP_403_FORBIDDEN,
+        )
+
+
+async def common_checks(  # noqa: PLR0915
     request_body: dict,
     team_object: Optional[LiteLLM_TeamTable],
     user_object: Optional[LiteLLM_UserTable],
@@ -477,6 +522,14 @@ async def common_checks(
         valid_token=valid_token,
     )
 
+    # 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path)
+    await check_tools_allowlist(
+        request_body=request_body,
+        valid_token=valid_token,
+        team_object=team_object,
+        route=route,
+    )
+
     return True
 
 
diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 0c25424ceaa..4c96e079c9e 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -13,49 +13,36 @@ import random
 import time
 import traceback
 from datetime import datetime, timedelta, timezone
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    Union,
-    cast,
-    overload,
-)
+from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union,
+                    cast, overload)
 
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import DualCache, RedisCache
 from litellm.constants import DB_SPEND_UPDATE_JOB_NAME
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
-from litellm.proxy._types import (
-    DB_CONNECTION_ERROR_TYPES,
-    BaseDailySpendTransaction,
-    DailyAgentSpendTransaction,
-    DailyEndUserSpendTransaction,
-    DailyOrganizationSpendTransaction,
-    DailyTagSpendTransaction,
-    DailyTeamSpendTransaction,
-    DailyUserSpendTransaction,
-    DBSpendUpdateTransactions,
-    Litellm_EntityType,
-    LiteLLM_UserTable,
-    SpendLogsMetadata,
-    SpendLogsPayload,
-    SpendUpdateQueueItem,
-    ToolDiscoveryQueueItem,
-)
-from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
-    DailySpendUpdateQueue,
-)
-from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
-from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
-from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
-from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
-    ToolDiscoveryQueue,
-)
+from litellm.proxy._types import (DB_CONNECTION_ERROR_TYPES,
+                                  BaseDailySpendTransaction,
+                                  DailyAgentSpendTransaction,
+                                  DailyEndUserSpendTransaction,
+                                  DailyOrganizationSpendTransaction,
+                                  DailyTagSpendTransaction,
+                                  DailyTeamSpendTransaction,
+                                  DailyUserSpendTransaction,
+                                  DBSpendUpdateTransactions,
+                                  Litellm_EntityType, LiteLLM_UserTable,
+                                  SpendLogsMetadata, SpendLogsPayload,
+                                  SpendUpdateQueueItem, ToolDiscoveryQueueItem)
+from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \
+    DailySpendUpdateQueue
+from litellm.proxy.db.db_transaction_queue.pod_lock_manager import \
+    PodLockManager
+from litellm.proxy.db.db_transaction_queue.redis_update_buffer import \
+    RedisUpdateBuffer
+from litellm.proxy.db.db_transaction_queue.spend_update_queue import \
+    SpendUpdateQueue
+from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import \
+    ToolDiscoveryQueue
 from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
 
 if TYPE_CHECKING:
@@ -104,12 +91,10 @@ class DBSpendUpdateWriter:
         end_time: Optional[datetime],
         response_cost: Optional[float],
     ):
-        from litellm.proxy.proxy_server import (
-            disable_spend_logs,
-            litellm_proxy_budget_name,
-            prisma_client,
-            user_api_key_cache,
-        )
+        from litellm.proxy.proxy_server import (disable_spend_logs,
+                                                litellm_proxy_budget_name,
+                                                prisma_client,
+                                                user_api_key_cache)
         from litellm.proxy.utils import ProxyUpdateSpend, hash_token
 
         try:
@@ -124,9 +109,8 @@ class DBSpendUpdateWriter:
                 hashed_token = token
 
             ## CREATE SPEND LOG PAYLOAD ##
-            from litellm.proxy.spend_tracking.spend_tracking_utils import (
-                get_logging_payload,
-            )
+            from litellm.proxy.spend_tracking.spend_tracking_utils import \
+                get_logging_payload
 
             payload = get_logging_payload(
                 kwargs=kwargs,
@@ -230,6 +214,7 @@ class DBSpendUpdateWriter:
             _litellm_params = kwargs.get("litellm_params") or {}
             _metadata = _litellm_params.get("metadata") or {}
             key_alias = _metadata.get("user_api_key_alias") or None
+            user_agent = _metadata.get("user_agent") or None
 
             def _enqueue(tool_name: str, origin: str = "user_defined") -> None:
                 self.tool_discovery_queue.add_update(
@@ -239,17 +224,20 @@ class DBSpendUpdateWriter:
                         key_hash=hashed_token,
                         team_id=team_id,
                         key_alias=key_alias,
+                        user_agent=user_agent,
                     )
                 )
 
             # --- MCP tool calls ---
             sl_object = kwargs.get("standard_logging_object")
             if sl_object is not None:
-                mcp_metadata = (
-                    sl_object.get("metadata", {}) or {}
-                ).get("mcp_tool_call_metadata")
+                mcp_metadata = (sl_object.get("metadata", {}) or {}).get(
+                    "mcp_tool_call_metadata"
+                )
                 if mcp_metadata and isinstance(mcp_metadata, dict):
-                    tool_name = mcp_metadata.get("namespaced_tool_name") or mcp_metadata.get("name")
+                    tool_name = mcp_metadata.get(
+                        "namespaced_tool_name"
+                    ) or mcp_metadata.get("name")
                     mcp_server_name = mcp_metadata.get("mcp_server_name")
                     if tool_name:
                         _enqueue(tool_name, origin=mcp_server_name or "user_defined")
@@ -280,7 +268,9 @@ class DBSpendUpdateWriter:
                     _enqueue(name)
 
             # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) ---
-            if completion_response is not None and hasattr(completion_response, "choices"):
+            if completion_response is not None and hasattr(
+                completion_response, "choices"
+            ):
                 for choice in completion_response.choices or []:
                     message = getattr(choice, "message", None)
                     if message is None:
@@ -768,19 +758,46 @@ class DBSpendUpdateWriter:
                     daily_end_user_spend_update_transactions,
                     daily_agent_spend_update_transactions,
                     daily_tag_spend_update_transactions,
-                ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
+                ) = (
+                    await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
+                )
 
                 if db_spend_update_transactions is not None:
                     verbose_proxy_logger.info(
                         "Spend tracking - committing spend updates from Redis to DB: "
                         "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d",
-                        len(db_spend_update_transactions.get("key_list_transactions") or {}),
-                        len(db_spend_update_transactions.get("user_list_transactions") or {}),
-                        len(db_spend_update_transactions.get("team_list_transactions") or {}),
-                        len(db_spend_update_transactions.get("org_list_transactions") or {}),
-                        len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
-                        len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
-                        len(db_spend_update_transactions.get("tag_list_transactions") or {}),
+                        len(
+                            db_spend_update_transactions.get("key_list_transactions")
+                            or {}
+                        ),
+                        len(
+                            db_spend_update_transactions.get("user_list_transactions")
+                            or {}
+                        ),
+                        len(
+                            db_spend_update_transactions.get("team_list_transactions")
+                            or {}
+                        ),
+                        len(
+                            db_spend_update_transactions.get("org_list_transactions")
+                            or {}
+                        ),
+                        len(
+                            db_spend_update_transactions.get(
+                                "end_user_list_transactions"
+                            )
+                            or {}
+                        ),
+                        len(
+                            db_spend_update_transactions.get(
+                                "team_member_list_transactions"
+                            )
+                            or {}
+                        ),
+                        len(
+                            db_spend_update_transactions.get("tag_list_transactions")
+                            or {}
+                        ),
                     )
                     await self._commit_spend_updates_to_db(
                         prisma_client=prisma_client,
@@ -985,10 +1002,8 @@ class DBSpendUpdateWriter:
         Commits all the spend `UPDATE` transactions to the Database
 
         """
-        from litellm.proxy.utils import (
-            ProxyUpdateSpend,
-            _raise_failed_update_spend_exception,
-        )
+        from litellm.proxy.utils import (ProxyUpdateSpend,
+                                         _raise_failed_update_spend_exception)
 
         ### UPDATE USER TABLE ###
         user_list_transactions = db_spend_update_transactions["user_list_transactions"]
@@ -1523,14 +1538,14 @@ class DBSpendUpdateWriter:
 
                                 # Add cache-related fields if they exist
                                 if "cache_read_input_tokens" in transaction:
-                                    common_data[
-                                        "cache_read_input_tokens"
-                                    ] = transaction.get("cache_read_input_tokens", 0)
+                                    common_data["cache_read_input_tokens"] = (
+                                        transaction.get("cache_read_input_tokens", 0)
+                                    )
                                 if "cache_creation_input_tokens" in transaction:
-                                    common_data[
-                                        "cache_creation_input_tokens"
-                                    ] = transaction.get(
-                                        "cache_creation_input_tokens", 0
+                                    common_data["cache_creation_input_tokens"] = (
+                                        transaction.get(
+                                            "cache_creation_input_tokens", 0
+                                        )
                                     )
 
                                 if entity_type == "tag" and "request_id" in transaction:
diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py
new file mode 100644
index 00000000000..6e8c63675e6
--- /dev/null
+++ b/litellm/proxy/db/spend_log_tool_index.py
@@ -0,0 +1,147 @@
+"""
+Track tool usage for the dashboard: insert into SpendLogToolIndex when spend logs
+are written, so "last N requests for tool X" and "how is this tool called in production"
+queries are fast.
+"""
+
+from datetime import datetime, timezone
+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
+
+
+def _add_tool_calls_to_set(tool_calls: Any, out: Set[str]) -> None:
+    """Extract tool names from OpenAI-style tool_calls list into out."""
+    if not isinstance(tool_calls, list):
+        return
+    for tc in tool_calls:
+        if not isinstance(tc, dict):
+            continue
+        fn = tc.get("function")
+        if isinstance(fn, dict):
+            name = fn.get("name")
+            if name and isinstance(name, str) and name.strip():
+                out.add(name.strip())
+
+
+def _parse_tool_names_from_payload(payload: Dict[str, Any]) -> Set[str]:
+    """
+    Extract deduplicated tool names from a spend log payload.
+    Sources: mcp_namespaced_tool_name, response (tool_calls), proxy_server_request (tools).
+    """
+    tool_names: Set[str] = set()
+
+    # Top-level MCP tool name (single tool per request for that flow)
+    mcp_name = payload.get("mcp_namespaced_tool_name")
+    if mcp_name and isinstance(mcp_name, str) and mcp_name.strip():
+        tool_names.add(mcp_name.strip())
+
+    # Response: OpenAI-style tool_calls[].function.name or choices[0].message.tool_calls
+    response_raw = payload.get("response")
+    if response_raw:
+        response_obj = (
+            safe_json_loads(response_raw, default=None)
+            if isinstance(response_raw, str)
+            else response_raw
+        )
+        if isinstance(response_obj, dict):
+            _add_tool_calls_to_set(response_obj.get("tool_calls"), tool_names)
+            choices = response_obj.get("choices")
+            if isinstance(choices, list) and choices:
+                msg = choices[0].get("message") if isinstance(choices[0], dict) else None
+                if isinstance(msg, dict):
+                    _add_tool_calls_to_set(msg.get("tool_calls"), tool_names)
+
+    # Request body: tools[].function.name
+    request_raw = payload.get("proxy_server_request")
+    if request_raw:
+        request_obj = (
+            safe_json_loads(request_raw, default=None)
+            if isinstance(request_raw, str)
+            else request_raw
+        )
+        if isinstance(request_obj, dict):
+            body = request_obj.get("body", request_obj)
+            if isinstance(body, dict):
+                request_obj = body
+        if isinstance(request_obj, dict):
+            tools = request_obj.get("tools")
+            if isinstance(tools, list):
+                for t in tools:
+                    if isinstance(t, dict):
+                        fn = t.get("function")
+                        if isinstance(fn, dict):
+                            name = fn.get("name")
+                            if name and isinstance(name, str) and name.strip():
+                                tool_names.add(name.strip())
+
+    return tool_names
+
+
+async def process_spend_logs_tool_usage(
+    prisma_client: PrismaClient,
+    logs_to_process: List[Dict[str, Any]],
+) -> None:
+    """
+    After spend logs are written: insert SpendLogToolIndex rows from each payload.
+    Extracts tool names from mcp_namespaced_tool_name, response tool_calls, and
+    proxy_server_request tools.
+    """
+    if not logs_to_process:
+        return
+
+    index_rows: List[Dict[str, Any]] = []
+
+    for payload in logs_to_process:
+        request_id = payload.get("request_id")
+        start_time = payload.get("startTime")
+        if not request_id or not start_time:
+            continue
+        if isinstance(start_time, str):
+            try:
+                start_time = datetime.fromisoformat(
+                    start_time.replace("Z", "+00:00")
+                )
+            except (ValueError, TypeError):
+                continue
+        if start_time.tzinfo is None:
+            start_time = start_time.replace(tzinfo=timezone.utc)
+
+        tool_names = _parse_tool_names_from_payload(payload)
+        for tool_name in tool_names:
+            index_rows.append({
+                "request_id": request_id,
+                "tool_name": tool_name,
+                "start_time": start_time,
+            })
+
+    if not index_rows:
+        return
+
+    try:
+        index_data = []
+        for r in index_rows:
+            st = r["start_time"]
+            if isinstance(st, str):
+                try:
+                    st = datetime.fromisoformat(st.replace("Z", "+00:00"))
+                except (ValueError, TypeError):
+                    continue
+            if st.tzinfo is None:
+                st = st.replace(tzinfo=timezone.utc)
+            index_data.append({
+                "request_id": r["request_id"],
+                "tool_name": r["tool_name"],
+                "start_time": st,
+            })
+        if index_data:
+            await prisma_client.db.litellm_spendlogtoolindex.create_many(
+                data=index_data,
+                skip_duplicates=True,
+            )
+    except Exception as e:
+        verbose_proxy_logger.warning(
+            "Tool usage tracking (SpendLogToolIndex) failed (non-fatal): %s", e
+        )
diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py
index 4e0a8095a08..0eda012d515 100644
--- a/litellm/proxy/db/tool_registry_writer.py
+++ b/litellm/proxy/db/tool_registry_writer.py
@@ -2,36 +2,64 @@
 DB helpers for LiteLLM_ToolTable — the global tool registry.
 
 Tools are auto-discovered from LLM responses and upserted here.
-Admins use the management endpoints to read and update call_policy.
-
-NOTE: Uses raw SQL (query_raw / execute_raw) instead of Prisma model methods
-because the generated Prisma Python client may not have LiteLLM_ToolTable
-when running against an older generated schema.
+Admins use the management endpoints to read and update input_policy / output_policy.
 """
 
 import uuid
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Dict, List, Optional
+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.types.tool_management import LiteLLM_ToolTableRow, ToolCallPolicy
+from litellm.types.tool_management import (
+    LiteLLM_ToolTableRow,
+    ToolPolicyOverrideRow,
+)
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient
 
 
-def _row_to_model(row: dict) -> LiteLLM_ToolTableRow:
+def _row_to_model(row: Union[dict, Any]) -> LiteLLM_ToolTableRow:
+    """Convert a Prisma model instance or dict to LiteLLM_ToolTableRow."""
+    model_dump = getattr(row, "model_dump", None)
+    if callable(model_dump):
+        row = model_dump()
+    elif not isinstance(row, dict):
+        row = {
+            k: getattr(row, k, None)
+            for k in (
+                "tool_id",
+                "tool_name",
+                "origin",
+                "input_policy",
+                "output_policy",
+                "call_count",
+                "assignments",
+                "key_hash",
+                "team_id",
+                "key_alias",
+                "user_agent",
+                "last_used_at",
+                "created_at",
+                "updated_at",
+                "created_by",
+                "updated_by",
+            )
+        }
     return LiteLLM_ToolTableRow(
         tool_id=row.get("tool_id", ""),
         tool_name=row.get("tool_name", ""),
         origin=row.get("origin"),
-        call_policy=row.get("call_policy", "untrusted"),
+        input_policy=row.get("input_policy") or "untrusted",
+        output_policy=row.get("output_policy") or "untrusted",
         call_count=int(row.get("call_count") or 0),
         assignments=row.get("assignments"),
         key_hash=row.get("key_hash"),
         team_id=row.get("team_id"),
         key_alias=row.get("key_alias"),
+        user_agent=row.get("user_agent"),
+        last_used_at=row.get("last_used_at"),
         created_at=row.get("created_at"),
         updated_at=row.get("updated_at"),
         created_by=row.get("created_by"),
@@ -44,10 +72,10 @@ async def batch_upsert_tools(
     items: List[ToolDiscoveryQueueItem],
 ) -> None:
     """
-    Batch-upsert tool registry rows via raw SQL.
+    Batch-upsert tool registry rows via Prisma.
 
-    On first insert: sets call_policy = "untrusted" (schema default), call_count = 1.
-    On conflict: increments call_count; preserves existing call_policy.
+    On first insert: sets input_policy/output_policy = "untrusted" (default), call_count = 1.
+    On conflict: increments call_count; preserves existing policies.
     """
     if not items:
         return
@@ -55,6 +83,8 @@ async def batch_upsert_tools(
         data = [item for item in items if item.get("tool_name")]
         if not data:
             return
+        now = datetime.now(timezone.utc)
+        table = prisma_client.db.litellm_tooltable
         for item in data:
             tool_name = item.get("tool_name", "")
             origin = item.get("origin") or "user_defined"
@@ -62,49 +92,52 @@ async def batch_upsert_tools(
             key_hash = item.get("key_hash")
             team_id = item.get("team_id")
             key_alias = item.get("key_alias")
-            now = datetime.now(timezone.utc).isoformat()
-            await prisma_client.db.execute_raw(
-                'INSERT INTO "LiteLLM_ToolTable" '
-                "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias, created_at, updated_at) "
-                "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8, $8) "
-                "ON CONFLICT (tool_name) DO UPDATE SET "
-                "call_count = \"LiteLLM_ToolTable\".call_count + 1, "
-                "updated_at = $8",
-                tool_name,
-                origin,
-                created_by,
-                key_hash,
-                team_id,
-                key_alias,
-                str(uuid.uuid4()),
-                now,
+            user_agent = item.get("user_agent")
+            await table.upsert(
+                where={"tool_name": tool_name},
+                data={
+                    "create": {
+                        "tool_id": str(uuid.uuid4()),
+                        "tool_name": tool_name,
+                        "origin": origin,
+                        "input_policy": "untrusted",
+                        "output_policy": "untrusted",
+                        "call_count": 1,
+                        "created_by": created_by,
+                        "updated_by": created_by,
+                        "key_hash": key_hash,
+                        "team_id": team_id,
+                        "key_alias": key_alias,
+                        "user_agent": user_agent,
+                        "last_used_at": now,
+                    },
+                    "update": {
+                        "call_count": {"increment": 1},
+                        "updated_at": now,
+                        "last_used_at": now,
+                    },
+                },
             )
         verbose_proxy_logger.debug(
             "tool_registry_writer: upserted %d tool(s)", len(data)
         )
     except Exception as e:
-        verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e)
+        verbose_proxy_logger.error(
+            "tool_registry_writer batch_upsert_tools error: %s", e
+        )
 
 
 async def list_tools(
     prisma_client: "PrismaClient",
-    call_policy: Optional[ToolCallPolicy] = None,
+    input_policy: Optional[str] = None,
 ) -> List[LiteLLM_ToolTableRow]:
-    """Return all tools, optionally filtered by call_policy."""
+    """Return all tools, optionally filtered by input_policy."""
     try:
-        if call_policy is not None:
-            rows = await prisma_client.db.query_raw(
-                'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, '
-                'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by '
-                'FROM "LiteLLM_ToolTable" WHERE call_policy = $1 ORDER BY created_at DESC',
-                call_policy,
-            )
-        else:
-            rows = await prisma_client.db.query_raw(
-                'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, '
-                'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by '
-                'FROM "LiteLLM_ToolTable" ORDER BY created_at DESC',
-            )
+        where = {"input_policy": input_policy} if input_policy is not None else {}
+        rows = await prisma_client.db.litellm_tooltable.find_many(
+            where=where,
+            order={"created_at": "desc"},
+        )
         return [_row_to_model(row) for row in rows]
     except Exception as e:
         verbose_proxy_logger.error("tool_registry_writer list_tools error: %s", e)
@@ -117,15 +150,12 @@ async def get_tool(
 ) -> Optional[LiteLLM_ToolTableRow]:
     """Return a single tool row by tool_name."""
     try:
-        rows = await prisma_client.db.query_raw(
-            'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, '
-            'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by '
-            'FROM "LiteLLM_ToolTable" WHERE tool_name = $1',
-            tool_name,
+        row = await prisma_client.db.litellm_tooltable.find_unique(
+            where={"tool_name": tool_name},
         )
-        if not rows:
+        if row is None:
             return None
-        return _row_to_model(rows[0])
+        return _row_to_model(row)
     except Exception as e:
         verbose_proxy_logger.error("tool_registry_writer get_tool error: %s", e)
         return None
@@ -134,46 +164,279 @@ async def get_tool(
 async def update_tool_policy(
     prisma_client: "PrismaClient",
     tool_name: str,
-    call_policy: ToolCallPolicy,
     updated_by: Optional[str],
+    input_policy: Optional[str] = None,
+    output_policy: Optional[str] = None,
 ) -> Optional[LiteLLM_ToolTableRow]:
-    """Update the call_policy for a tool. Upserts the row if it does not exist yet."""
+    """Update input_policy and/or output_policy for a tool. Upserts the row if it does not exist yet."""
     try:
         _updated_by = updated_by or "system"
-        now = datetime.now(timezone.utc).isoformat()
-        await prisma_client.db.execute_raw(
-            'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by, created_at, updated_at) '
-            "VALUES ($4, $1, $2, $3, $3, $5, $5) "
-            "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5",
-            tool_name,
-            call_policy,
-            _updated_by,
-            str(uuid.uuid4()),
-            now,
+        now = datetime.now(timezone.utc)
+
+        create_data: dict = {
+            "tool_id": str(uuid.uuid4()),
+            "tool_name": tool_name,
+            "input_policy": input_policy or "untrusted",
+            "output_policy": output_policy or "untrusted",
+            "created_by": _updated_by,
+            "updated_by": _updated_by,
+            "created_at": now,
+            "updated_at": now,
+        }
+        update_data: dict = {
+            "updated_by": _updated_by,
+            "updated_at": now,
+        }
+        if input_policy is not None:
+            update_data["input_policy"] = input_policy
+        if output_policy is not None:
+            update_data["output_policy"] = output_policy
+
+        await prisma_client.db.litellm_tooltable.upsert(
+            where={"tool_name": tool_name},
+            data={
+                "create": create_data,
+                "update": update_data,
+            },
         )
         return await get_tool(prisma_client, tool_name)
     except Exception as e:
-        verbose_proxy_logger.error("tool_registry_writer update_tool_policy error: %s", e)
+        verbose_proxy_logger.error(
+            "tool_registry_writer update_tool_policy error: %s", e
+        )
         return None
 
 
 async def get_tools_by_names(
     prisma_client: "PrismaClient",
     tool_names: List[str],
-) -> Dict[str, str]:
+) -> Dict[str, Tuple[str, str]]:
     """
-    Return a {tool_name: call_policy} map for the given tool names.
-    Used by the policy enforcement guardrail — single batch query, never N+1.
+    Return a {tool_name: (input_policy, output_policy)} map for the given tool names.
     """
     if not tool_names:
         return {}
     try:
-        placeholders = ", ".join(f"${i+1}" for i in range(len(tool_names)))
-        rows = await prisma_client.db.query_raw(
-            f'SELECT tool_name, call_policy FROM "LiteLLM_ToolTable" WHERE tool_name IN ({placeholders})',
-            *tool_names,
+        rows = await prisma_client.db.litellm_tooltable.find_many(
+            where={"tool_name": {"in": tool_names}},
         )
-        return {row["tool_name"]: row["call_policy"] for row in rows}
+        return {
+            row.tool_name: (
+                getattr(row, "input_policy", "untrusted") or "untrusted",
+                getattr(row, "output_policy", "untrusted") or "untrusted",
+            )
+            for row in rows
+        }
     except Exception as e:
-        verbose_proxy_logger.error("tool_registry_writer get_tools_by_names error: %s", e)
+        verbose_proxy_logger.error(
+            "tool_registry_writer get_tools_by_names error: %s", e
+        )
         return {}
+
+
+async def list_overrides_for_tool(
+    prisma_client: "PrismaClient",
+    tool_name: str,
+) -> List[ToolPolicyOverrideRow]:
+    """
+    Return override-like rows for a tool by finding object permissions that have
+    this tool in blocked_tools, then resolving each permission to key/team scope for display.
+    """
+    out: List[ToolPolicyOverrideRow] = []
+    try:
+        perms = await prisma_client.db.litellm_objectpermissiontable.find_many(
+            where={"blocked_tools": {"has": tool_name}},
+            include={
+                "verification_tokens": True,
+                "teams": True,
+            },
+        )
+        for perm in perms:
+            op_id = getattr(perm, "object_permission_id", None) or ""
+            tokens = getattr(perm, "verification_tokens", []) or []
+            teams = getattr(perm, "teams", []) or []
+            for t in tokens:
+                out.append(
+                    ToolPolicyOverrideRow(
+                        override_id=op_id,
+                        tool_name=tool_name,
+                        team_id=None,
+                        key_hash=getattr(t, "token", None),
+                        input_policy="blocked",
+                        key_alias=getattr(t, "key_alias", None),
+                        created_at=None,
+                        updated_at=None,
+                    )
+                )
+            for team in teams:
+                out.append(
+                    ToolPolicyOverrideRow(
+                        override_id=op_id,
+                        tool_name=tool_name,
+                        team_id=getattr(team, "team_id", None),
+                        key_hash=None,
+                        input_policy="blocked",
+                        key_alias=getattr(team, "team_alias", None),
+                        created_at=None,
+                        updated_at=None,
+                    )
+                )
+        return out
+    except Exception as e:
+        verbose_proxy_logger.error(
+            "tool_registry_writer list_overrides_for_tool error: %s", e
+        )
+        return []
+
+
+class ToolPolicyRegistry:
+    """
+    In-memory registry of tool policies synced from DB.
+    Hot path uses get_effective_policies only — no DB, no cache.
+    """
+
+    def __init__(self) -> None:
+        self._tool_input_policies: Dict[str, str] = {}
+        self._tool_output_policies: Dict[str, str] = {}
+        self._blocked_tools_by_op_id: Dict[str, List[str]] = {}
+        self._initialized: bool = False
+
+    def is_initialized(self) -> bool:
+        return self._initialized
+
+    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()
+            self._tool_input_policies = {
+                row.tool_name: getattr(row, "input_policy", "untrusted") or "untrusted"
+                for row in tools
+            }
+            self._tool_output_policies = {
+                row.tool_name: getattr(row, "output_policy", "untrusted") or "untrusted"
+                for row in tools
+            }
+
+            perms = await prisma_client.db.litellm_objectpermissiontable.find_many()
+            self._blocked_tools_by_op_id = {}
+            for row in perms:
+                op_id = getattr(row, "object_permission_id", None)
+                blocked = getattr(row, "blocked_tools", None) or []
+                if op_id:
+                    self._blocked_tools_by_op_id[op_id] = list(blocked)
+
+            self._initialized = True
+            verbose_proxy_logger.info(
+                "ToolPolicyRegistry: synced %d tool policies and %d object permissions from DB",
+                len(self._tool_input_policies),
+                len(self._blocked_tools_by_op_id),
+            )
+        except Exception as e:
+            verbose_proxy_logger.exception(
+                "ToolPolicyRegistry sync_tool_policy_from_db error: %s", e
+            )
+            raise
+
+    def get_input_policy(self, tool_name: str) -> str:
+        return self._tool_input_policies.get(tool_name, "untrusted")
+
+    def get_output_policy(self, tool_name: str) -> str:
+        return self._tool_output_policies.get(tool_name, "untrusted")
+
+    def get_effective_policies(
+        self,
+        tool_names: List[str],
+        object_permission_id: Optional[str] = None,
+        team_object_permission_id: Optional[str] = None,
+    ) -> Dict[str, str]:
+        """
+        Return effective input_policy per tool from in-memory state.
+        If tool is in key or team blocked_tools -> "blocked", else global input_policy or "untrusted".
+        """
+        if not tool_names:
+            return {}
+        blocked: set = set()
+        for op_id in (object_permission_id, team_object_permission_id):
+            if op_id and op_id.strip():
+                blocked.update(
+                    self._blocked_tools_by_op_id.get(op_id.strip(), [])
+                )
+        result: Dict[str, str] = {}
+        for name in tool_names:
+            if name in blocked:
+                result[name] = "blocked"
+            else:
+                result[name] = self._tool_input_policies.get(name, "untrusted")
+        return result
+
+
+_tool_policy_registry: Optional[ToolPolicyRegistry] = None
+
+
+def get_tool_policy_registry() -> ToolPolicyRegistry:
+    """Return the global ToolPolicyRegistry singleton."""
+    global _tool_policy_registry
+    if _tool_policy_registry is None:
+        _tool_policy_registry = ToolPolicyRegistry()
+    return _tool_policy_registry
+
+
+async def add_tool_to_object_permission_blocked(
+    prisma_client: "PrismaClient",
+    object_permission_id: str,
+    tool_name: str,
+) -> bool:
+    """Add tool_name to the permission's blocked_tools if not already present."""
+    if not object_permission_id or not tool_name:
+        return False
+    try:
+        row = await prisma_client.db.litellm_objectpermissiontable.find_unique(
+            where={"object_permission_id": object_permission_id},
+        )
+        if row is None:
+            return False
+        current = list(getattr(row, "blocked_tools", []) or [])
+        if tool_name in current:
+            return True
+        current.append(tool_name)
+        await prisma_client.db.litellm_objectpermissiontable.update(
+            where={"object_permission_id": object_permission_id},
+            data={"blocked_tools": current},
+        )
+        return True
+    except Exception as e:
+        verbose_proxy_logger.error(
+            "tool_registry_writer add_tool_to_object_permission_blocked error: %s", e
+        )
+        return False
+
+
+async def remove_tool_from_object_permission_blocked(
+    prisma_client: "PrismaClient",
+    object_permission_id: str,
+    tool_name: str,
+) -> bool:
+    """Remove tool_name from the permission's blocked_tools. Returns False if tool was not in list."""
+    if not object_permission_id or not tool_name:
+        return False
+    try:
+        row = await prisma_client.db.litellm_objectpermissiontable.find_unique(
+            where={"object_permission_id": object_permission_id},
+        )
+        if row is None:
+            return False
+        current = list(getattr(row, "blocked_tools", []) or [])
+        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(
+            where={"object_permission_id": object_permission_id},
+            data={"blocked_tools": current},
+        )
+        return True
+    except Exception as e:
+        verbose_proxy_logger.error(
+            "tool_registry_writer remove_tool_from_object_permission_blocked error: %s",
+            e,
+        )
+        return False
diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py
index 87558566c42..368948414e9 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py
@@ -1,13 +1,16 @@
 """
 Tool Policy Guardrail
 
-Reads call_policy from LiteLLM_ToolTable and enforces it on LLM requests/responses.
+Reads input_policy / output_policy from LiteLLM_ToolTable and enforces them.
 
-Policy values:
-  "trusted"   - allow through (no action)
-  "untrusted" - allow through (no action; default for newly discovered tools)
+Input policy values:
+  "untrusted" - allow through (default for newly discovered tools)
+  "trusted"   - only allow if conversation contains no untrusted tool output
   "blocked"   - raise HTTPException, preventing the tool call
-  "dual_llm"  - (Phase 3) send to second LLM for verification; currently treated as allowed
+
+Output policy values:
+  "untrusted" - output may be tainted (default)
+  "trusted"   - output is verified safe
 
 Configuration in proxy config YAML:
   guardrails:
@@ -15,25 +18,18 @@ Configuration in proxy config YAML:
       litellm_params:
         guardrail: tool_policy
         mode: post_call
-
-or both pre and post call:
-    - guardrail_name: "tool_policy"
-      litellm_params:
-        guardrail: tool_policy
-        mode: during_call  # runs before LLM and on response
 """
 
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple
 
 from fastapi import HTTPException
 
 from litellm._logging import verbose_proxy_logger
-from litellm.caching.dual_cache import DualCache
-from litellm.constants import TOOL_POLICY_CACHE_TTL_SECONDS
 from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
     log_guardrail_information,
 )
+from litellm.proxy.guardrails.tool_name_extraction import extract_request_tool_names
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import GenericGuardrailAPIInputs
 
@@ -43,12 +39,71 @@ if TYPE_CHECKING:
 GUARDRAIL_NAME = "tool_policy"
 
 
+def _get_request_object_permission_ids(
+    request_data: dict,
+) -> Tuple[Optional[str], Optional[str]]:
+    """Extract object_permission_id and team_object_permission_id from request_data."""
+    if not request_data:
+        return None, None
+    for key in ("litellm_metadata", "metadata"):
+        meta = request_data.get(key)
+        if not isinstance(meta, dict):
+            continue
+        auth = meta.get("user_api_key_auth")
+        if auth is not None and hasattr(auth, "object_permission_id"):
+            key_op = getattr(auth, "object_permission_id", None)
+            team_op = getattr(auth, "team_object_permission_id", None)
+            if key_op is not None or team_op is not None:
+                return (
+                    str(key_op).strip() if key_op else None,
+                    str(team_op).strip() if team_op else None,
+                )
+        key_op = meta.get("user_api_key_object_permission_id")
+        team_op = meta.get("user_api_key_team_object_permission_id")
+        if key_op is not None or team_op is not None:
+            return (
+                str(key_op).strip() if key_op else None,
+                str(team_op).strip() if team_op else None,
+            )
+    return None, None
+
+
+def _get_request_route_from_data(request_data: dict) -> Optional[str]:
+    """Get request route from request_data (metadata or top-level)."""
+    route = request_data.get("user_api_key_request_route")
+    if route:
+        return route
+    meta = request_data.get("metadata") or request_data.get("litellm_metadata") or {}
+    return meta.get("user_api_key_request_route")
+
+
+def _resolve_tool_names_from_messages(messages: List[dict]) -> Dict[str, str]:
+    """
+    Build a map of tool_call_id -> tool_name from assistant messages' tool_calls.
+    Used to resolve which tool produced each tool result in the conversation.
+    """
+    mapping: Dict[str, str] = {}
+    for msg in messages:
+        if msg.get("role") != "assistant":
+            continue
+        tool_calls = msg.get("tool_calls") or []
+        for tc in tool_calls:
+            if isinstance(tc, dict):
+                tc_id = tc.get("id")
+                fn = (tc.get("function") or {}).get("name")
+            else:
+                tc_id = getattr(tc, "id", None)
+                fn_obj = getattr(tc, "function", None)
+                fn = getattr(fn_obj, "name", None) if fn_obj else None
+            if tc_id and fn:
+                mapping[tc_id] = fn
+    return mapping
+
+
 class ToolPolicyGuardrail(CustomGuardrail):
     """
-    Guardrail that enforces per-tool call policies stored in LiteLLM_ToolTable.
-
-    Tools with call_policy="blocked" are rejected before/after the LLM call.
-    Tools with call_policy="trusted" or "untrusted" pass through unchanged.
+    Guardrail that enforces per-tool input/output policies from the in-memory
+    ToolPolicyRegistry (synced from DB).
     """
 
     def __init__(self, **kwargs: Any) -> None:
@@ -59,7 +114,6 @@ class ToolPolicyGuardrail(CustomGuardrail):
                 GuardrailEventHooks.during_call,
             ]
         super().__init__(**kwargs)
-        self._policy_cache: DualCache = DualCache()
 
     @log_guardrail_information
     async def apply_guardrail(
@@ -70,12 +124,7 @@ class ToolPolicyGuardrail(CustomGuardrail):
         logging_obj: Optional["LiteLLMLoggingObj"] = None,
     ) -> GenericGuardrailAPIInputs:
         """
-        Enforce tool policies on both request tools and response tool_calls.
-
-        - input_type="request":  check inputs["tools"] (tool definitions in the LLM request)
-        - input_type="response": check inputs["tool_calls"] (tool_calls in the LLM response)
-
-        Raises HTTPException (400) if any tool is "blocked".
+        Enforce input_policy and output_policy trust chain on request tools / response tool_calls.
         """
         if input_type == "request":
             tools = inputs.get("tools") or []
@@ -86,7 +135,11 @@ class ToolPolicyGuardrail(CustomGuardrail):
                 and isinstance(t.get("function"), dict)
                 and t["function"].get("name")
             ]
-        else:  # response
+            if not tool_names:
+                route = _get_request_route_from_data(request_data)
+                if route:
+                    tool_names = extract_request_tool_names(route, request_data)
+        else:
             tool_calls = inputs.get("tool_calls") or []
             tool_names = []
             for tc in tool_calls:
@@ -101,12 +154,25 @@ class ToolPolicyGuardrail(CustomGuardrail):
         if not tool_names:
             return inputs
 
-        policy_map = await self._get_policies_cached(tool_names)
+        object_permission_id, team_object_permission_id = (
+            _get_request_object_permission_ids(request_data)
+        )
+        from litellm.proxy.db.tool_registry_writer import get_tool_policy_registry
 
+        registry = get_tool_policy_registry()
+        if not registry.is_initialized():
+            return inputs
+
+        # Stage 1: Check for blocked tools (input_policy=blocked or per-key/team override)
+        policy_map = registry.get_effective_policies(
+            tool_names,
+            object_permission_id=object_permission_id,
+            team_object_permission_id=team_object_permission_id,
+        )
         blocked = [name for name in tool_names if policy_map.get(name) == "blocked"]
         if blocked:
             verbose_proxy_logger.warning(
-                "ToolPolicyGuardrail: blocking tool(s) %s (policy=blocked)", blocked
+                "ToolPolicyGuardrail: blocking tool(s) %s (input_policy=blocked)", blocked
             )
             raise HTTPException(
                 status_code=400,
@@ -117,47 +183,47 @@ class ToolPolicyGuardrail(CustomGuardrail):
                 },
             )
 
+        # Stage 2: Trust chain enforcement (response path only)
+        # For each tool with input_policy=trusted, check if conversation
+        # contains output from tools with output_policy=untrusted
+        if input_type == "response":
+            trusted_input_tools = [
+                name for name in tool_names if policy_map.get(name) == "trusted"
+            ]
+            if trusted_input_tools:
+                messages = request_data.get("messages") or []
+                tc_id_to_name = _resolve_tool_names_from_messages(messages)
+
+                untrusted_sources: List[str] = []
+                for msg in messages:
+                    if msg.get("role") != "tool":
+                        continue
+                    tool_call_id = msg.get("tool_call_id")
+                    source_tool = tc_id_to_name.get(tool_call_id, "") if tool_call_id else ""
+                    if not source_tool:
+                        continue
+                    if registry.get_output_policy(source_tool) == "untrusted":
+                        if source_tool not in untrusted_sources:
+                            untrusted_sources.append(source_tool)
+
+                if untrusted_sources:
+                    verbose_proxy_logger.warning(
+                        "ToolPolicyGuardrail: trust chain violation — %s require trusted input "
+                        "but conversation has untrusted output from %s",
+                        trusted_input_tools,
+                        untrusted_sources,
+                    )
+                    raise HTTPException(
+                        status_code=400,
+                        detail={
+                            "error": "Violated tool policy",
+                            "blocked_tools": trusted_input_tools,
+                            "untrusted_sources": untrusted_sources,
+                            "message": (
+                                f"{', '.join(trusted_input_tools)} requires trusted input but "
+                                f"conversation contains untrusted output from {', '.join(untrusted_sources)}."
+                            ),
+                        },
+                    )
+
         return inputs
-
-    async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]:
-        """
-        Batch-fetch call_policy for the given tool names.
-
-        Caches per individual tool name (not per combination) so that adding
-        a new tool to a request doesn't invalidate the cached policies for all
-        the other tools already in the cache.
-        """
-        from litellm.proxy.db.tool_registry_writer import get_tools_by_names
-        from litellm.proxy.proxy_server import prisma_client
-
-        if not tool_names or prisma_client is None:
-            return {}
-
-        result: Dict[str, str] = {}
-        cache_misses: List[str] = []
-
-        for name in tool_names:
-            cached = await self._policy_cache.async_get_cache(f"tool_policy:{name}")
-            if cached is not None and isinstance(cached, str):
-                result[name] = cached
-            else:
-                cache_misses.append(name)
-
-        if cache_misses:
-            fetched = await get_tools_by_names(
-                prisma_client=prisma_client, tool_names=cache_misses
-            )
-            for name, policy in fetched.items():
-                result[name] = policy
-                await self._policy_cache.async_set_cache(
-                    key=f"tool_policy:{name}",
-                    value=policy,
-                    ttl=TOOL_POLICY_CACHE_TTL_SECONDS,
-                )
-            verbose_proxy_logger.debug(
-                "ToolPolicyGuardrail: fetched %d policies from DB (cache hits: %d)",
-                len(cache_misses),
-                len(tool_names) - len(cache_misses),
-            )
-
-        return result
diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py
new file mode 100644
index 00000000000..db24fa2277c
--- /dev/null
+++ b/litellm/proxy/guardrails/tool_name_extraction.py
@@ -0,0 +1,85 @@
+"""
+Extract tool names from request body by route/call type.
+
+Used by auth (check_tools_allowlist) and ToolPolicyGuardrail so tool-format
+knowledge lives in one place. Uses guardrail translation handlers where available,
+with standalone extractors for generate_content and MCP.
+"""
+
+from typing import Any, Dict, List
+
+from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
+from litellm.llms import load_guardrail_translation_mappings
+from litellm.types.utils import CallTypes
+
+# Call types that have no guardrail translation handler; we use standalone extractors
+STANDALONE_EXTRACTORS: Dict[str, Any] = {}
+
+
+def _extract_generate_content_tool_names(data: dict) -> List[str]:
+    """Google generateContent: tools[].functionDeclarations[].name"""
+    names: List[str] = []
+    for tool in data.get("tools") or []:
+        if not isinstance(tool, dict):
+            continue
+        for decl in tool.get("functionDeclarations") or []:
+            if isinstance(decl, dict) and decl.get("name"):
+                names.append(str(decl["name"]))
+    return names
+
+
+def _extract_mcp_tool_names(data: dict) -> List[str]:
+    """MCP call_tool: name or mcp_tool_name in body"""
+    names: List[str] = []
+    name = data.get("name") or data.get("mcp_tool_name")
+    if name:
+        names.append(str(name))
+    return names
+
+
+def _register_standalone_extractors() -> None:
+    if STANDALONE_EXTRACTORS:
+        return
+    STANDALONE_EXTRACTORS[CallTypes.generate_content.value] = _extract_generate_content_tool_names
+    STANDALONE_EXTRACTORS[CallTypes.agenerate_content.value] = _extract_generate_content_tool_names
+    STANDALONE_EXTRACTORS[CallTypes.call_mcp_tool.value] = _extract_mcp_tool_names
+
+
+# Tool-capable call types (routes that can send tools in the request)
+TOOL_CAPABLE_CALL_TYPES = frozenset({
+    CallTypes.completion.value,
+    CallTypes.acompletion.value,
+    CallTypes.responses.value,
+    CallTypes.aresponses.value,
+    CallTypes.anthropic_messages.value,
+    CallTypes.generate_content.value,
+    CallTypes.agenerate_content.value,
+    CallTypes.call_mcp_tool.value,
+})
+
+
+def extract_request_tool_names(route: str, data: dict) -> List[str]:
+    """
+    Extract tool names from the request body for the given route.
+    Uses guardrail translation handlers when available, else standalone extractors
+    for generate_content and MCP. Returns [] for non-tool-capable routes or when
+    no tools are present.
+    """
+    call_types = get_call_types_for_route(route)
+    if not call_types:
+        return []
+    _register_standalone_extractors()
+    mappings = load_guardrail_translation_mappings()
+    for call_type in call_types:
+        if not isinstance(call_type, CallTypes):
+            continue
+        if call_type.value not in TOOL_CAPABLE_CALL_TYPES:
+            continue
+        if call_type.value in STANDALONE_EXTRACTORS:
+            return STANDALONE_EXTRACTORS[call_type.value](data)
+        handler_cls = mappings.get(call_type)
+        if handler_cls is not None:
+            names = handler_cls().extract_request_tool_names(data)
+            if names:
+                return names
+    return []
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index 68f1ab114bf..32eab99fb99 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -1091,6 +1091,15 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
     ] = user_api_key_dict.user_max_budget
 
     data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata
+    data[_metadata_variable_name]["user_api_key_team_metadata"] = (
+        user_api_key_dict.team_metadata
+    )
+    data[_metadata_variable_name]["user_api_key_object_permission_id"] = (
+        getattr(user_api_key_dict, "object_permission_id", None)
+    )
+    data[_metadata_variable_name]["user_api_key_team_object_permission_id"] = (
+        getattr(user_api_key_dict, "team_object_permission_id", None)
+    )
     data[_metadata_variable_name]["headers"] = _headers
     data[_metadata_variable_name]["endpoint"] = str(request.url)
 
diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py
index 89880c9a4ec..7fdd3475c04 100644
--- a/litellm/proxy/management_endpoints/tool_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py
@@ -4,27 +4,87 @@ TOOL POLICY MANAGEMENT
 All /tool management endpoints
 
 GET  /v1/tool/list              - List all discovered tools and their policies
+GET  /v1/tool/policy/options    - List available input/output policy options with descriptions
 GET  /v1/tool/{tool_name}       - Get a single tool's details
-POST /v1/tool/policy            - Update the call_policy for a tool
+POST /v1/tool/policy            - Update the input_policy / output_policy for a tool
 """
 
-from typing import Optional
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, List, Optional
 
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Depends, HTTPException, Query
+
+if TYPE_CHECKING:
+    from litellm.proxy.utils import PrismaClient
 
 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.types.tool_management import (
     LiteLLM_ToolTableRow,
-    ToolCallPolicy,
+    ToolDetailResponse,
+    ToolInputPolicy,
     ToolListResponse,
+    ToolOutputPolicy,
+    ToolPolicyOption,
+    ToolPolicyOptionsResponse,
     ToolPolicyUpdateRequest,
     ToolPolicyUpdateResponse,
+    ToolUsageLogEntry,
+    ToolUsageLogsResponse,
 )
 
 router = APIRouter()
 
+TOOL_POLICY_OPTIONS = ToolPolicyOptionsResponse(
+    input_policies=[
+        ToolPolicyOption(
+            value="untrusted",
+            label="Untrusted",
+            description="Tool accepts any input, including data from untrusted tool outputs. Default for newly discovered tools.",
+        ),
+        ToolPolicyOption(
+            value="trusted",
+            label="Trusted",
+            description="Tool requires trusted input. Blocked if the conversation contains output from any tool with output_policy=untrusted.",
+        ),
+        ToolPolicyOption(
+            value="blocked",
+            label="Blocked",
+            description="Tool is completely prohibited. Any attempt to call it is rejected.",
+        ),
+    ],
+    output_policies=[
+        ToolPolicyOption(
+            value="untrusted",
+            label="Untrusted",
+            description="Tool output may contain unsafe content (prompt injection, risky code). Downstream tools with input_policy=trusted will be blocked.",
+        ),
+        ToolPolicyOption(
+            value="trusted",
+            label="Trusted",
+            description="Tool output is verified safe. Will not trigger trust-chain blocks on downstream tools.",
+        ),
+    ],
+)
+
+
+@router.get(
+    "/v1/tool/policy/options",
+    tags=["tool management"],
+    dependencies=[Depends(user_api_key_auth)],
+    response_model=ToolPolicyOptionsResponse,
+)
+async def get_tool_policy_options(
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Return the available input and output policy options with descriptions.
+    Static data — no DB call.
+    """
+    return TOOL_POLICY_OPTIONS
+
 
 @router.get(
     "/v1/tool/list",
@@ -33,14 +93,14 @@ router = APIRouter()
     response_model=ToolListResponse,
 )
 async def list_tools(
-    call_policy: Optional[ToolCallPolicy] = None,
+    input_policy: Optional[ToolInputPolicy] = None,
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
-    List all auto-discovered tools and their call policies.
+    List all auto-discovered tools and their policies.
 
     Parameters:
-    - call_policy: Optional filter — one of "trusted", "untrusted", "dual_llm", "blocked"
+    - input_policy: Optional filter — one of "trusted", "untrusted", "blocked"
     """
     from litellm.proxy.db.tool_registry_writer import list_tools as db_list_tools
     from litellm.proxy.proxy_server import prisma_client
@@ -51,13 +111,201 @@ async def list_tools(
         )
 
     try:
-        tools = await db_list_tools(prisma_client=prisma_client, call_policy=call_policy)
+        tools = await db_list_tools(
+            prisma_client=prisma_client, input_policy=input_policy
+        )
         return ToolListResponse(tools=tools, total=len(tools))
     except Exception as e:
         verbose_proxy_logger.exception("Error listing tools: %s", e)
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@router.get(
+    "/v1/tool/{tool_name:path}/detail",
+    tags=["tool management"],
+    dependencies=[Depends(user_api_key_auth)],
+    response_model=ToolDetailResponse,
+)
+async def get_tool_detail(
+    tool_name: str,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Get a single tool with its policy overrides (for UI detail view).
+    """
+    from litellm.proxy.db.tool_registry_writer import get_tool as db_get_tool
+    from litellm.proxy.db.tool_registry_writer import list_overrides_for_tool
+    from litellm.proxy.proxy_server import prisma_client
+
+    if prisma_client is None:
+        raise HTTPException(
+            status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
+        )
+
+    try:
+        tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name)
+        if tool is None:
+            raise HTTPException(status_code=404, detail=f"Tool '{tool_name}' not found")
+        overrides = await list_overrides_for_tool(
+            prisma_client=prisma_client, tool_name=tool_name
+        )
+        return ToolDetailResponse(tool=tool, overrides=overrides)
+    except HTTPException:
+        raise
+    except Exception as e:
+        verbose_proxy_logger.exception("Error getting tool detail: %s", e)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> Optional[str]:
+    """Short snippet from messages or proxy_server_request for tool usage log row."""
+    if sl is None:
+        return None
+    messages = getattr(sl, "messages", None)
+    if messages is not None:
+        s = _snippet_str(messages, max_len)
+        if s:
+            return s
+    psr = getattr(sl, "proxy_server_request", None)
+    if not psr:
+        return None
+    if isinstance(psr, str):
+        import json
+
+        try:
+            psr = json.loads(psr)
+        except Exception:
+            return _snippet_str(psr, max_len)
+    if isinstance(psr, dict):
+        msgs = psr.get("messages")
+        if msgs is None and isinstance(psr.get("body"), dict):
+            msgs = psr["body"].get("messages")
+        s = _snippet_str(msgs, max_len)
+        if s:
+            return s
+    return _snippet_str(psr, max_len)
+
+
+def _snippet_str(text: Any, max_len: int = 200) -> Optional[str]:
+    if text is None:
+        return None
+    if isinstance(text, str):
+        s = text
+    elif isinstance(text, list):
+        parts = []
+        for item in text:
+            if isinstance(item, dict) and "content" in item:
+                c = item["content"]
+                parts.append(c if isinstance(c, str) else str(c))
+            else:
+                parts.append(str(item))
+        s = " ".join(parts)
+    else:
+        s = str(text)
+    if not s or s == "{}":
+        return None
+    return (s[:max_len] + "...") if len(s) > max_len else s
+
+
+@router.get(
+    "/v1/tool/{tool_name:path}/logs",
+    tags=["tool management"],
+    dependencies=[Depends(user_api_key_auth)],
+    response_model=ToolUsageLogsResponse,
+)
+async def get_tool_usage_logs(
+    tool_name: str,
+    page: int = Query(1, ge=1),
+    page_size: int = Query(50, ge=1, le=100),
+    start_date: Optional[str] = Query(None, description="YYYY-MM-DD"),
+    end_date: Optional[str] = Query(None, description="YYYY-MM-DD"),
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Return paginated spend logs for requests that used this tool (from SpendLogToolIndex).
+    """
+    from litellm.proxy.proxy_server import prisma_client
+
+    if prisma_client is None:
+        raise HTTPException(
+            status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
+        )
+
+    try:
+        where: dict = {"tool_name": tool_name}
+        if start_date or end_date:
+            start_time_filter: Optional[datetime] = None
+            end_time_filter: Optional[datetime] = None
+            if start_date:
+                try:
+                    start_time_filter = datetime.strptime(
+                        start_date + "T00:00:00", "%Y-%m-%dT%H:%M:%S"
+                    ).replace(tzinfo=timezone.utc)
+                except ValueError:
+                    pass
+            if end_date:
+                try:
+                    end_time_filter = datetime.strptime(
+                        end_date + "T23:59:59", "%Y-%m-%dT%H:%M:%S"
+                    ).replace(tzinfo=timezone.utc)
+                except ValueError:
+                    pass
+            if start_time_filter is not None or end_time_filter is not None:
+                where["start_time"] = {}
+                if start_time_filter is not None:
+                    where["start_time"]["gte"] = start_time_filter
+                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(
+            where=where,
+            order={"start_time": "desc"},
+            skip=(page - 1) * page_size,
+            take=page_size,
+        )
+        request_ids = [r.request_id for r in index_rows]
+        if not request_ids:
+            return ToolUsageLogsResponse(
+                logs=[], total=total, page=page, page_size=page_size
+            )
+
+        spend_logs = await prisma_client.db.litellm_spendlogs.find_many(
+            where={"request_id": {"in": request_ids}}
+        )
+        log_by_id = {s.request_id: s for s in spend_logs}
+
+        logs_out: List[ToolUsageLogEntry] = []
+        for r in index_rows:
+            sl = log_by_id.get(r.request_id)
+            if not sl:
+                continue
+            ts = (
+                sl.startTime.isoformat()
+                if hasattr(sl.startTime, "isoformat")
+                else str(sl.startTime)
+            )
+            logs_out.append(
+                ToolUsageLogEntry(
+                    id=sl.request_id,
+                    timestamp=ts,
+                    model=getattr(sl, "model", None) or None,
+                    spend=getattr(sl, "spend", None),
+                    total_tokens=getattr(sl, "total_tokens", None),
+                    input_snippet=_input_snippet_for_tool_log(sl),
+                )
+            )
+
+        return ToolUsageLogsResponse(
+            logs=logs_out, total=total, page=page, page_size=page_size
+        )
+    except HTTPException:
+        raise
+    except Exception as e:
+        verbose_proxy_logger.exception("Error getting tool usage logs: %s", e)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @router.get(
     "/v1/tool/{tool_name:path}",
     tags=["tool management"],
@@ -70,9 +318,6 @@ async def get_tool(
 ):
     """
     Get details for a single tool.
-
-    Parameters:
-    - tool_name: The tool name (supports namespaced names with slashes)
     """
     from litellm.proxy.db.tool_registry_writer import get_tool as db_get_tool
     from litellm.proxy.proxy_server import prisma_client
@@ -85,9 +330,7 @@ async def get_tool(
     try:
         tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name)
         if tool is None:
-            raise HTTPException(
-                status_code=404, detail=f"Tool '{tool_name}' not found"
-            )
+            raise HTTPException(status_code=404, detail=f"Tool '{tool_name}' not found")
         return tool
     except HTTPException:
         raise
@@ -96,6 +339,80 @@ async def get_tool(
         raise HTTPException(status_code=500, detail=str(e))
 
 
+async def _resolve_key_hash_to_object_permission_id(
+    prisma_client: "PrismaClient",
+    key_hash: str,
+) -> Optional[str]:
+    """Resolve key (hash or raw) to object_permission_id; create permission if key has none."""
+    from litellm.proxy.proxy_server import hash_token
+
+    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(
+        where={"token": hashed}
+    )
+    if row is None:
+        return None
+    op_id = getattr(row, "object_permission_id", None)
+    if op_id:
+        return op_id
+    new_id = str(uuid.uuid4())
+    await prisma_client.db.litellm_objectpermissiontable.create(
+        data={"object_permission_id": new_id, "blocked_tools": []}
+    )
+    updated_count = await prisma_client.db.litellm_verificationtoken.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(
+            where={"object_permission_id": new_id}
+        )
+        row = await prisma_client.db.litellm_verificationtoken.find_unique(
+            where={"token": hashed}
+        )
+        return getattr(row, "object_permission_id", None) if row else None
+    return new_id
+
+
+async def _resolve_team_id_to_object_permission_id(
+    prisma_client: "PrismaClient",
+    team_id: str,
+) -> Optional[str]:
+    """Resolve team_id to object_permission_id; create permission if team has none."""
+    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(
+        where={"team_id": team_id_clean},
+        select={"object_permission_id": True},
+    )
+    if row is None:
+        return None
+    op_id = getattr(row, "object_permission_id", None)
+    if op_id:
+        return op_id
+    new_id = str(uuid.uuid4())
+    await prisma_client.db.litellm_objectpermissiontable.create(
+        data={"object_permission_id": new_id, "blocked_tools": []}
+    )
+    updated_count = await prisma_client.db.litellm_teamtable.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(
+            where={"object_permission_id": new_id}
+        )
+        row = await prisma_client.db.litellm_teamtable.find_unique(
+            where={"team_id": team_id_clean},
+            select={"object_permission_id": True},
+        )
+        return getattr(row, "object_permission_id", None) if row else None
+    return new_id
+
+
 @router.post(
     "/v1/tool/policy",
     tags=["tool management"],
@@ -107,15 +424,20 @@ async def update_tool_policy(
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
-    Set the call policy for a tool.
+    Set the input_policy and/or output_policy for a tool (global), or block for a specific team/key (override).
 
     Parameters:
     - tool_name: str - The tool to update
-    - call_policy: "trusted" | "untrusted" | "dual_llm" | "blocked"
-
-    Setting a tool to "blocked" will cause the ToolPolicyGuardrail to remove
-    that tool_call from LLM responses before returning them to the client.
+    - input_policy: optional - "trusted" | "untrusted" | "blocked"
+    - output_policy: optional - "trusted" | "untrusted"
+    - team_id: optional - if set, create/update override for this team only
+    - key_hash: optional - if set, create/update override for this key only
     """
+    from litellm.proxy.db.tool_registry_writer import (
+        add_tool_to_object_permission_blocked,
+        get_tool_policy_registry,
+        remove_tool_from_object_permission_blocked,
+    )
     from litellm.proxy.db.tool_registry_writer import (
         update_tool_policy as db_update_tool_policy,
     )
@@ -127,19 +449,80 @@ async def update_tool_policy(
         )
 
     try:
+        if data.team_id is not None or data.key_hash is not None:
+            if data.team_id is not None and data.key_hash is not None:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Provide either team_id or key_hash, not both",
+                )
+            if data.key_hash is not None:
+                op_id = await _resolve_key_hash_to_object_permission_id(
+                    prisma_client, data.key_hash
+                )
+            else:
+                op_id = await _resolve_team_id_to_object_permission_id(
+                    prisma_client, data.team_id or ""
+                )
+            if op_id is None:
+                raise HTTPException(
+                    status_code=404,
+                    detail="Key or team not found for the given identifier",
+                )
+            is_blocking = data.input_policy == "blocked"
+            if is_blocking:
+                ok = await add_tool_to_object_permission_blocked(
+                    prisma_client=prisma_client,
+                    object_permission_id=op_id,
+                    tool_name=data.tool_name,
+                )
+            else:
+                ok = await remove_tool_from_object_permission_blocked(
+                    prisma_client=prisma_client,
+                    object_permission_id=op_id,
+                    tool_name=data.tool_name,
+                )
+            if not ok:
+                raise HTTPException(
+                    status_code=500,
+                    detail=f"Failed to update policy override for tool '{data.tool_name}'",
+                )
+            registry = get_tool_policy_registry()
+            if registry.is_initialized():
+                await registry.sync_tool_policy_from_db(prisma_client)
+            return ToolPolicyUpdateResponse(
+                tool_name=data.tool_name,
+                input_policy=data.input_policy,
+                output_policy=data.output_policy,
+                updated=True,
+                team_id=data.team_id,
+                key_hash=data.key_hash,
+            )
+
+        if data.input_policy is None and data.output_policy is None:
+            raise HTTPException(
+                status_code=400,
+                detail="At least one of input_policy or output_policy must be provided",
+            )
+
         updated = await db_update_tool_policy(
             prisma_client=prisma_client,
             tool_name=data.tool_name,
-            call_policy=data.call_policy,
             updated_by=user_api_key_dict.user_id,
+            input_policy=data.input_policy,
+            output_policy=data.output_policy,
         )
         if updated is None:
             raise HTTPException(
-                status_code=500, detail=f"Failed to update policy for tool '{data.tool_name}'"
+                status_code=500,
+                detail=f"Failed to update policy for tool '{data.tool_name}'",
             )
+        registry = get_tool_policy_registry()
+        if registry.is_initialized():
+            await registry.sync_tool_policy_from_db(prisma_client)
         return ToolPolicyUpdateResponse(
             tool_name=updated.tool_name,
-            call_policy=updated.call_policy,
+            input_policy=updated.input_policy,
+            output_policy=updated.output_policy,
             updated=True,
         )
     except HTTPException:
@@ -147,3 +530,77 @@ async def update_tool_policy(
     except Exception as e:
         verbose_proxy_logger.exception("Error updating tool policy: %s", e)
         raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.delete(
+    "/v1/tool/{tool_name:path}/overrides",
+    tags=["tool management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def delete_tool_policy_override(
+    tool_name: str,
+    team_id: Optional[str] = Query(
+        None, description="Team ID of the override to remove"
+    ),
+    key_hash: Optional[str] = Query(
+        None, description="Key hash of the override to remove"
+    ),
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Remove a policy override for a tool. Specify the override by team_id or key_hash
+    (exactly one required).
+    """
+    from litellm.proxy.db.tool_registry_writer import (
+        get_tool_policy_registry,
+        remove_tool_from_object_permission_blocked,
+    )
+    from litellm.proxy.proxy_server import prisma_client
+
+    if prisma_client is None:
+        raise HTTPException(
+            status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
+        )
+    if team_id is None and key_hash is None:
+        raise HTTPException(
+            status_code=400,
+            detail="At least one of team_id or key_hash is required to identify the override",
+        )
+    if team_id is not None and key_hash is not None:
+        raise HTTPException(
+            status_code=400,
+            detail="Provide either team_id or key_hash, not both",
+        )
+    try:
+        if key_hash is not None:
+            op_id = await _resolve_key_hash_to_object_permission_id(
+                prisma_client, key_hash
+            )
+        else:
+            op_id = await _resolve_team_id_to_object_permission_id(
+                prisma_client, team_id or ""
+            )
+        if op_id is None:
+            raise HTTPException(
+                status_code=404,
+                detail="Key or team not found for the given identifier",
+            )
+        deleted = await remove_tool_from_object_permission_blocked(
+            prisma_client=prisma_client,
+            object_permission_id=op_id,
+            tool_name=tool_name,
+        )
+        if not deleted:
+            raise HTTPException(
+                status_code=404,
+                detail=f"No override found for tool '{tool_name}' with the given scope",
+            )
+        registry = get_tool_policy_registry()
+        if registry.is_initialized():
+            await registry.sync_tool_policy_from_db(prisma_client)
+        return {"deleted": True, "tool_name": tool_name}
+    except HTTPException:
+        raise
+    except Exception as e:
+        verbose_proxy_logger.exception("Error deleting tool policy override: %s", e)
+        raise HTTPException(status_code=500, detail=str(e))
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index b3d707b1aa2..6a2b0accb0e 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -4411,6 +4411,9 @@ class ProxyConfig:
         if self._should_load_db_object(object_type="search_tools"):
             await self._init_search_tools_in_db(prisma_client=prisma_client)
 
+        if self._should_load_db_object(object_type="tools"):
+            await self._init_tool_policy_in_db(prisma_client=prisma_client)
+
         if self._should_load_db_object(object_type="model_cost_map"):
             await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
 
@@ -4847,6 +4850,24 @@ class ProxyConfig:
                 )
             )
 
+    async def _init_tool_policy_in_db(self, prisma_client: PrismaClient):
+        """
+        Initialize tool policy from database into the in-memory registry.
+        Synced periodically by add_deployment -> _init_non_llm_objects_in_db.
+        """
+        from litellm.proxy.db.tool_registry_writer import get_tool_policy_registry
+
+        try:
+            registry = get_tool_policy_registry()
+            await registry.sync_tool_policy_from_db(prisma_client=prisma_client)
+            verbose_proxy_logger.debug("Successfully synced tool policy from DB")
+        except Exception as e:
+            verbose_proxy_logger.exception(
+                "litellm.proxy.proxy_server.py::ProxyConfig:_init_tool_policy_in_db - {}".format(
+                    str(e)
+                )
+            )
+
     async def _init_vector_stores_in_db(self, prisma_client: PrismaClient):
         from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
 
@@ -10577,6 +10598,12 @@ async def async_queue_request(
         data["metadata"]["user_api_key_team_id"] = getattr(
             user_api_key_dict, "team_id", None
         )
+        data["metadata"]["user_api_key_object_permission_id"] = getattr(
+            user_api_key_dict, "object_permission_id", None
+        )
+        data["metadata"]["user_api_key_team_object_permission_id"] = getattr(
+            user_api_key_dict, "team_object_permission_id", None
+        )
         data["metadata"]["endpoint"] = str(request.url)
 
         global user_temperature, user_request_timeout, user_max_tokens, user_api_base
@@ -11093,9 +11120,7 @@ async def get_favicon():
 
     if favicon_url.startswith(("http://", "https://")):
         try:
-            from litellm.llms.custom_httpx.http_handler import (
-                get_async_httpx_client,
-            )
+            from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
             from litellm.types.llms.custom_http import httpxSpecialProvider
 
             async_client = get_async_httpx_client(
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index e0b28a4e012..25ee2750548 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -260,6 +260,7 @@ model LiteLLM_ObjectPermissionTable {
   vector_stores         String[]       @default([])
   agents                String[]       @default([])
   agent_access_groups   String[]       @default([])
+  blocked_tools         String[]       @default([]) // Tool names blocked for any key/team/user with this permission
   teams                 LiteLLM_TeamTable[]
   projects              LiteLLM_ProjectTable[]
   verification_tokens   LiteLLM_VerificationToken[]
@@ -928,6 +929,16 @@ model LiteLLM_SpendLogGuardrailIndex {
   @@index([policy_id, start_time])
 }
 
+// Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production
+model LiteLLM_SpendLogToolIndex {
+  request_id String
+  tool_name  String   // matches LiteLLM_ToolTable.tool_name; join for input_policy/output_policy etc.
+  start_time DateTime
+
+  @@id([request_id, tool_name])
+  @@index([tool_name, start_time])
+}
+
 // Prompt table for storing prompt configurations
 model LiteLLM_PromptTable {
   id String @id @default(uuid())
@@ -1065,23 +1076,27 @@ model LiteLLM_PolicyAttachmentTable {
   updated_by       String?
 }
 
-// Global tool registry - auto-discovered from LLM responses; admins set call_policy here
+// Global tool registry - auto-discovered from LLM responses; admins set input/output policies here
 model LiteLLM_ToolTable {
-  tool_id     String   @id @default(uuid())
-  tool_name   String   @unique   // e.g. "huggingface_remote-mcp__dynamic_space"
-  origin      String?            // MCP server name or "user_defined"
-  call_policy String   @default("untrusted")  // "trusted" | "untrusted" | "dual_llm" | "blocked"
-  call_count  Int      @default(0)            // cumulative number of times this tool was seen
-  assignments Json?    @default("{}")
-  key_hash    String?            // hash of the virtual key that first called this tool
-  team_id     String?            // team that first called this tool
-  key_alias   String?            // human-readable alias of the virtual key
-  created_at  DateTime @default(now())
-  created_by  String?
-  updated_at  DateTime @default(now()) @updatedAt
-  updated_by  String?
+  tool_id        String   @id @default(uuid())
+  tool_name      String   @unique   // e.g. "huggingface_remote-mcp__dynamic_space"
+  origin         String?            // MCP server name or "user_defined"
+  input_policy   String   @default("untrusted")  // "trusted" | "untrusted" | "blocked"
+  output_policy  String   @default("untrusted")  // "trusted" | "untrusted"
+  call_count     Int      @default(0)            // cumulative number of times this tool was seen
+  assignments    Json?    @default("{}")
+  key_hash       String?            // hash of the virtual key that first called this tool
+  team_id        String?            // team that first called this tool
+  key_alias      String?            // human-readable alias of the virtual key
+  user_agent     String?            // user-agent of the first request that discovered this tool
+  last_used_at   DateTime?          // timestamp of the most recent call
+  created_at     DateTime @default(now())
+  created_by     String?
+  updated_at     DateTime @default(now()) @updatedAt
+  updated_by     String?
 
-  @@index([call_policy])
+  @@index([input_policy])
+  @@index([output_policy])
   @@index([team_id])
 }
 
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index afcdd9d0c50..e6da95bb78f 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -3583,8 +3583,9 @@ class PrismaClient:
     def _get_engine_pid(self) -> int:
         try:
             engine = self.db._original_prisma._engine  # type: ignore[attr-defined]
-            if engine is not None and engine.process is not None:
-                return engine.process.pid
+            process = getattr(engine, "process", None) if engine is not None else None
+            if process is not None:
+                return process.pid
         except (AttributeError, TypeError):
             pass
         return 0
@@ -4688,6 +4689,19 @@ async def update_spend_logs_job(
             guardrail_tracking_err,
         )
 
+    # Tool usage tracking (same batch): SpendLogToolIndex for "last N requests for tool X"
+    try:
+        from litellm.proxy.db.spend_log_tool_index import process_spend_logs_tool_usage
+        await process_spend_logs_tool_usage(
+            prisma_client=prisma_client,
+            logs_to_process=logs_to_process,
+        )
+    except Exception as tool_tracking_err:
+        verbose_proxy_logger.warning(
+            "Spend tracking - tool usage tracking failed (non-fatal): %s",
+            tool_tracking_err,
+        )
+
 
 async def _monitor_spend_logs_queue(
     prisma_client: PrismaClient,
diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py
index 8704ff27759..1c5e1df9e9a 100644
--- a/litellm/types/tool_management.py
+++ b/litellm/types/tool_management.py
@@ -5,21 +5,27 @@ Pydantic models for Tool Policy management endpoints.
 from datetime import datetime
 from typing import Dict, List, Literal, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 ToolCallPolicy = Literal["trusted", "untrusted", "dual_llm", "blocked"]
 
+ToolInputPolicy = Literal["trusted", "untrusted", "blocked"]
+ToolOutputPolicy = Literal["trusted", "untrusted"]
+
 
 class LiteLLM_ToolTableRow(BaseModel):
     tool_id: str
     tool_name: str
     origin: Optional[str] = None
-    call_policy: ToolCallPolicy = "untrusted"
+    input_policy: ToolInputPolicy = "untrusted"
+    output_policy: ToolOutputPolicy = "untrusted"
     call_count: int = 0
     assignments: Optional[Dict] = None
     key_hash: Optional[str] = None
     team_id: Optional[str] = None
     key_alias: Optional[str] = None
+    user_agent: Optional[str] = None
+    last_used_at: Optional[datetime] = None
     created_at: Optional[datetime] = None
     updated_at: Optional[datetime] = None
     created_by: Optional[str] = None
@@ -33,10 +39,62 @@ class ToolListResponse(BaseModel):
 
 class ToolPolicyUpdateRequest(BaseModel):
     tool_name: str
-    call_policy: ToolCallPolicy
+    input_policy: Optional[ToolInputPolicy] = None
+    output_policy: Optional[ToolOutputPolicy] = None
+    team_id: Optional[str] = None
+    key_hash: Optional[str] = None
+    key_alias: Optional[str] = None
 
 
 class ToolPolicyUpdateResponse(BaseModel):
     tool_name: str
-    call_policy: ToolCallPolicy
+    input_policy: Optional[ToolInputPolicy] = None
+    output_policy: Optional[ToolOutputPolicy] = None
     updated: bool
+    team_id: Optional[str] = None
+    key_hash: Optional[str] = None
+
+
+class ToolPolicyOverrideRow(BaseModel):
+    override_id: str
+    tool_name: str
+    team_id: Optional[str] = None
+    key_hash: Optional[str] = None
+    input_policy: ToolInputPolicy = "blocked"
+    key_alias: Optional[str] = None
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+
+
+class ToolPolicyOption(BaseModel):
+    value: str
+    label: str
+    description: str
+
+
+class ToolPolicyOptionsResponse(BaseModel):
+    input_policies: List[ToolPolicyOption]
+    output_policies: List[ToolPolicyOption]
+
+
+class ToolDetailResponse(BaseModel):
+    tool: LiteLLM_ToolTableRow
+    overrides: List[ToolPolicyOverrideRow] = Field(default_factory=list)
+
+
+class ToolUsageLogEntry(BaseModel):
+    """One spend log row for a tool call (for UI "recent logs" table)."""
+
+    id: str  # request_id
+    timestamp: str
+    model: Optional[str] = None
+    spend: Optional[float] = None
+    total_tokens: Optional[int] = None
+    input_snippet: Optional[str] = None
+
+
+class ToolUsageLogsResponse(BaseModel):
+    logs: List[ToolUsageLogEntry]
+    total: int
+    page: int
+    page_size: int
diff --git a/schema.prisma b/schema.prisma
index e0b28a4e012..25ee2750548 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -260,6 +260,7 @@ model LiteLLM_ObjectPermissionTable {
   vector_stores         String[]       @default([])
   agents                String[]       @default([])
   agent_access_groups   String[]       @default([])
+  blocked_tools         String[]       @default([]) // Tool names blocked for any key/team/user with this permission
   teams                 LiteLLM_TeamTable[]
   projects              LiteLLM_ProjectTable[]
   verification_tokens   LiteLLM_VerificationToken[]
@@ -928,6 +929,16 @@ model LiteLLM_SpendLogGuardrailIndex {
   @@index([policy_id, start_time])
 }
 
+// Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production
+model LiteLLM_SpendLogToolIndex {
+  request_id String
+  tool_name  String   // matches LiteLLM_ToolTable.tool_name; join for input_policy/output_policy etc.
+  start_time DateTime
+
+  @@id([request_id, tool_name])
+  @@index([tool_name, start_time])
+}
+
 // Prompt table for storing prompt configurations
 model LiteLLM_PromptTable {
   id String @id @default(uuid())
@@ -1065,23 +1076,27 @@ model LiteLLM_PolicyAttachmentTable {
   updated_by       String?
 }
 
-// Global tool registry - auto-discovered from LLM responses; admins set call_policy here
+// Global tool registry - auto-discovered from LLM responses; admins set input/output policies here
 model LiteLLM_ToolTable {
-  tool_id     String   @id @default(uuid())
-  tool_name   String   @unique   // e.g. "huggingface_remote-mcp__dynamic_space"
-  origin      String?            // MCP server name or "user_defined"
-  call_policy String   @default("untrusted")  // "trusted" | "untrusted" | "dual_llm" | "blocked"
-  call_count  Int      @default(0)            // cumulative number of times this tool was seen
-  assignments Json?    @default("{}")
-  key_hash    String?            // hash of the virtual key that first called this tool
-  team_id     String?            // team that first called this tool
-  key_alias   String?            // human-readable alias of the virtual key
-  created_at  DateTime @default(now())
-  created_by  String?
-  updated_at  DateTime @default(now()) @updatedAt
-  updated_by  String?
+  tool_id        String   @id @default(uuid())
+  tool_name      String   @unique   // e.g. "huggingface_remote-mcp__dynamic_space"
+  origin         String?            // MCP server name or "user_defined"
+  input_policy   String   @default("untrusted")  // "trusted" | "untrusted" | "blocked"
+  output_policy  String   @default("untrusted")  // "trusted" | "untrusted"
+  call_count     Int      @default(0)            // cumulative number of times this tool was seen
+  assignments    Json?    @default("{}")
+  key_hash       String?            // hash of the virtual key that first called this tool
+  team_id        String?            // team that first called this tool
+  key_alias      String?            // human-readable alias of the virtual key
+  user_agent     String?            // user-agent of the first request that discovered this tool
+  last_used_at   DateTime?          // timestamp of the most recent call
+  created_at     DateTime @default(now())
+  created_by     String?
+  updated_at     DateTime @default(now()) @updatedAt
+  updated_by     String?
 
-  @@index([call_policy])
+  @@index([input_policy])
+  @@index([output_policy])
   @@index([team_id])
 }
 
diff --git a/scripts/test_tool_allowlist_script.py b/scripts/test_tool_allowlist_script.py
new file mode 100644
index 00000000000..75a50d09b84
--- /dev/null
+++ b/scripts/test_tool_allowlist_script.py
@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+"""
+Standalone script to test tool allowlist enforcement and tool name extraction.
+
+Run from repo root:
+  poetry run python scripts/test_tool_allowlist_script.py
+
+Or run the unit tests:
+  poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+# Ensure repo root is on path
+repo_root = Path(__file__).resolve().parent.parent
+if str(repo_root) not in sys.path:
+    sys.path.insert(0, str(repo_root))
+
+
+def test_extraction():
+    """Test extract_request_tool_names for each API shape."""
+    from litellm.proxy.guardrails.tool_name_extraction import extract_request_tool_names
+
+    cases = [
+        ("OpenAI chat tools", "/v1/chat/completions", {"tools": [{"type": "function", "function": {"name": "get_weather"}}]}),
+        ("OpenAI chat functions", "/v1/chat/completions", {"functions": [{"name": "run_sql"}]}),
+        ("OpenAI responses function", "/v1/responses", {"tools": [{"type": "function", "name": "get_current_weather"}]}),
+        ("OpenAI responses MCP", "/v1/responses", {"tools": [{"type": "mcp", "server_label": "dmcp"}]}),
+        ("Anthropic", "/v1/messages", {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]}),
+        ("Google generateContent", "/generate_content", {"tools": [{"functionDeclarations": [{"name": "schedule_meeting"}]}]}),
+        ("MCP call_tool", "/mcp/call_tool", {"name": "my_tool", "arguments": {}}),
+        ("Non-tool route", "/v1/embeddings", {"tools": [{"type": "function", "function": {"name": "x"}}]}),
+    ]
+    print("=== extract_request_tool_names(route, data) ===\n")
+    for label, route, data in cases:
+        names = extract_request_tool_names(route, data)
+        print(f"  {label}: {names}")
+    print()
+
+
+async def test_check_tools_allowlist():
+    """Test check_tools_allowlist with mock tokens."""
+    from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
+    from litellm.proxy.auth.auth_checks import check_tools_allowlist
+
+    def token(metadata=None, team_metadata=None):
+        return UserAPIKeyAuth(
+            api_key="test-key",
+            user_id="user",
+            team_id="team",
+            org_id=None,
+            models=["*"],
+            metadata=metadata or {},
+            team_metadata=team_metadata or {},
+        )
+
+    print("=== check_tools_allowlist (auth) ===\n")
+
+    # No allowlist -> pass
+    await check_tools_allowlist(
+        request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
+        valid_token=token(),
+        team_object=None,
+        route="/v1/chat/completions",
+    )
+    print("  No allowlist, body has tools: PASS")
+
+    # Allowed tool -> pass
+    await check_tools_allowlist(
+        request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
+        valid_token=token(metadata={"allowed_tools": ["get_weather"]}),
+        team_object=None,
+        route="/v1/chat/completions",
+    )
+    print("  allowed_tools=['get_weather'], body has get_weather: PASS")
+
+    # Disallowed tool -> raise
+    try:
+        await check_tools_allowlist(
+            request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
+            valid_token=token(metadata={"allowed_tools": ["other_tool"]}),
+            team_object=None,
+            route="/v1/chat/completions",
+        )
+        print("  DISALLOWED: expected ProxyException")
+    except ProxyException as e:
+        if e.type == ProxyErrorTypes.tool_access_denied:
+            print("  allowed_tools=['other_tool'], body has get_weather: PASS (raised tool_access_denied)")
+        else:
+            print(f"  Unexpected ProxyException type: {e.type}")
+    except Exception as e:
+        print(f"  Unexpected: {e}")
+
+    # Team allowlist when key empty
+    await check_tools_allowlist(
+        request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
+        valid_token=token(team_metadata={"allowed_tools": ["get_weather"]}),
+        team_object=None,
+        route="/v1/chat/completions",
+    )
+    print("  team_metadata.allowed_tools=['get_weather']: PASS")
+    print()
+
+
+def main():
+    print("Tool allowlist / tool name extraction – script checks\n")
+    test_extraction()
+    asyncio.run(test_check_tools_allowlist())
+    print("Done. For full unit tests run:")
+    print("  poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v")
+
+
+if __name__ == "__main__":
+    main()
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 44f9e32058a..1b1ee7afcba 100644
--- a/tests/test_litellm/proxy/db/test_tool_registry_writer.py
+++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py
@@ -1,6 +1,6 @@
 """
 Unit tests for tool_registry_writer.py — uses a mock prisma client
-that exposes execute_raw / query_raw (matching the actual raw-SQL implementation).
+that exposes litellm_tooltable.upsert / find_many / find_unique.
 """
 
 import os
@@ -13,21 +13,28 @@ import pytest
 sys.path.insert(0, os.path.abspath("../../.."))
 
 from litellm.proxy.db.tool_registry_writer import (
+    ToolPolicyRegistry,
     batch_upsert_tools,
     get_tool,
+    get_tool_policy_registry,
     get_tools_by_names,
     list_tools,
     update_tool_policy,
 )
 
 
-def _make_prisma(query_rows=None):
-    """Return a minimal mock prisma_client with execute_raw / query_raw."""
-    default_row = {
+def _mock_row(**kwargs):
+    """Build a row-like object with real attributes (no MagicMock) for _row_to_model."""
+
+    class Row:
+        pass
+
+    default = {
         "tool_id": "uuid-1",
         "tool_name": "my_tool",
         "origin": "user_defined",
-        "call_policy": "untrusted",
+        "input_policy": "untrusted",
+        "output_policy": "untrusted",
         "call_count": 1,
         "assignments": {},
         "key_hash": None,
@@ -38,31 +45,54 @@ def _make_prisma(query_rows=None):
         "created_by": None,
         "updated_by": None,
     }
-    rows = query_rows if query_rows is not None else [default_row]
+    default.update(kwargs)
+    row = Row()
+    for k, v in default.items():
+        setattr(row, k, v)
+    return row
 
+
+def _make_prisma(
+    *,
+    upsert_return=None,
+    find_many_rows=None,
+    find_unique_row=None,
+):
+    """Return a mock prisma_client with litellm_tooltable.upsert, find_many, find_unique."""
     prisma = MagicMock()
-    prisma.db.execute_raw = AsyncMock(return_value=None)
-    prisma.db.query_raw = AsyncMock(return_value=rows)
+    prisma.db.litellm_tooltable = MagicMock()
+    prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return)
+    prisma.db.litellm_tooltable.find_many = AsyncMock(
+        return_value=find_many_rows if find_many_rows is not None else []
+    )
+    prisma.db.litellm_tooltable.find_unique = AsyncMock(
+        return_value=find_unique_row
+    )
     return prisma
 
 
 @pytest.mark.asyncio
-async def test_batch_upsert_tools_calls_execute_raw():
+async def test_batch_upsert_tools_calls_upsert():
     prisma = _make_prisma()
     items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}]
     await batch_upsert_tools(prisma, items)
-    prisma.db.execute_raw.assert_awaited_once()
-    call_args = prisma.db.execute_raw.call_args
-    sql = call_args.args[0]
-    assert "LiteLLM_ToolTable" in sql
-    assert "ON CONFLICT" in sql
+    prisma.db.litellm_tooltable.upsert.assert_awaited_once()
+    call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs
+    assert call_kw["where"] == {"tool_name": "tool_a"}
+    assert call_kw["data"]["create"]["tool_name"] == "tool_a"
+    assert call_kw["data"]["create"]["origin"] == "mcp_server"
+    assert call_kw["data"]["create"]["input_policy"] == "untrusted"
+    assert call_kw["data"]["create"]["output_policy"] == "untrusted"
+    assert call_kw["data"]["create"]["call_count"] == 1
+    assert call_kw["data"]["update"]["call_count"] == {"increment": 1}
+    assert "updated_at" in call_kw["data"]["update"]
 
 
 @pytest.mark.asyncio
 async def test_batch_upsert_tools_empty_list():
     prisma = _make_prisma()
     await batch_upsert_tools(prisma, [])
-    prisma.db.execute_raw.assert_not_awaited()
+    prisma.db.litellm_tooltable.upsert.assert_not_awaited()
 
 
 @pytest.mark.asyncio
@@ -70,123 +100,120 @@ async def test_batch_upsert_tools_skips_empty_names():
     prisma = _make_prisma()
     items = [{"tool_name": "", "origin": None}, {"tool_name": None}]  # type: ignore[list-item]
     await batch_upsert_tools(prisma, items)
-    prisma.db.execute_raw.assert_not_awaited()
+    prisma.db.litellm_tooltable.upsert.assert_not_awaited()
 
 
 @pytest.mark.asyncio
-async def test_batch_upsert_multiple_tools_calls_execute_raw_per_tool():
+async def test_batch_upsert_multiple_tools_calls_upsert_per_tool():
     prisma = _make_prisma()
     items = [
         {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None},
         {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"},
     ]
     await batch_upsert_tools(prisma, items)
-    assert prisma.db.execute_raw.await_count == 2
+    assert prisma.db.litellm_tooltable.upsert.await_count == 2
+    calls = prisma.db.litellm_tooltable.upsert.call_args_list
+    assert calls[0].kwargs["where"]["tool_name"] == "tool_a"
+    assert calls[1].kwargs["where"]["tool_name"] == "tool_b"
 
 
 @pytest.mark.asyncio
 async def test_list_tools_no_filter():
-    row = {
-        "tool_id": "id1",
-        "tool_name": "tool_a",
-        "origin": "mcp",
-        "call_policy": "untrusted",
-        "call_count": 5,
-        "assignments": {},
-        "key_hash": None,
-        "team_id": None,
-        "key_alias": None,
-        "created_at": datetime.now(timezone.utc),
-        "updated_at": datetime.now(timezone.utc),
-        "created_by": None,
-        "updated_by": None,
-    }
-    prisma = _make_prisma(query_rows=[row])
+    row = _mock_row(
+        tool_id="id1",
+        tool_name="tool_a",
+        origin="mcp",
+        input_policy="untrusted",
+        output_policy="untrusted",
+        call_count=5,
+    )
+    prisma = _make_prisma(find_many_rows=[row])
     result = await list_tools(prisma)
     assert len(result) == 1
     assert result[0].tool_name == "tool_a"
     assert result[0].call_count == 5
-    prisma.db.query_raw.assert_awaited_once()
+    prisma.db.litellm_tooltable.find_many.assert_awaited_once()
+    call_kw = prisma.db.litellm_tooltable.find_many.call_args.kwargs
+    assert call_kw["where"] == {}
+    assert call_kw["order"] == {"created_at": "desc"}
 
 
 @pytest.mark.asyncio
-async def test_list_tools_with_policy_filter():
-    row = {
-        "tool_id": "id1",
-        "tool_name": "blocked_tool",
-        "origin": None,
-        "call_policy": "blocked",
-        "call_count": 2,
-        "assignments": None,
-        "key_hash": None,
-        "team_id": None,
-        "key_alias": None,
-        "created_at": datetime.now(timezone.utc),
-        "updated_at": datetime.now(timezone.utc),
-        "created_by": None,
-        "updated_by": None,
-    }
-    prisma = _make_prisma(query_rows=[row])
-    result = await list_tools(prisma, call_policy="blocked")
-    assert result[0].call_policy == "blocked"
-    call_args = prisma.db.query_raw.call_args
-    sql = call_args.args[0]
-    assert "WHERE call_policy" in sql
+async def test_list_tools_with_input_policy_filter():
+    row = _mock_row(
+        tool_id="id1",
+        tool_name="blocked_tool",
+        origin=None,
+        input_policy="blocked",
+        output_policy="untrusted",
+        call_count=2,
+        assignments=None,
+    )
+    prisma = _make_prisma(find_many_rows=[row])
+    result = await list_tools(prisma, input_policy="blocked")
+    assert result[0].input_policy == "blocked"
+    call_kw = prisma.db.litellm_tooltable.find_many.call_args.kwargs
+    assert call_kw["where"] == {"input_policy": "blocked"}
 
 
 @pytest.mark.asyncio
 async def test_get_tool_found():
-    prisma = _make_prisma()
+    row = _mock_row(tool_name="my_tool")
+    prisma = _make_prisma(find_unique_row=row)
     result = await get_tool(prisma, "my_tool")
     assert result is not None
     assert result.tool_name == "my_tool"
-    prisma.db.query_raw.assert_awaited_once()
+    prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(
+        where={"tool_name": "my_tool"}
+    )
 
 
 @pytest.mark.asyncio
 async def test_get_tool_not_found():
-    prisma = _make_prisma(query_rows=[])
+    prisma = _make_prisma(find_unique_row=None)
     result = await get_tool(prisma, "nonexistent")
     assert result is None
 
 
 @pytest.mark.asyncio
-async def test_update_tool_policy_calls_execute_raw():
-    row = {
-        "tool_id": "uuid-1",
-        "tool_name": "my_tool",
-        "origin": "user_defined",
-        "call_policy": "blocked",
-        "call_count": 1,
-        "assignments": {},
-        "key_hash": None,
-        "team_id": None,
-        "key_alias": None,
-        "created_at": datetime.now(timezone.utc),
-        "updated_at": datetime.now(timezone.utc),
-        "created_by": None,
-        "updated_by": "admin",
-    }
-    prisma = _make_prisma(query_rows=[row])
-    result = await update_tool_policy(prisma, "my_tool", "blocked", "admin")
+async def test_update_tool_policy_calls_upsert_then_get_tool():
+    row = _mock_row(
+        tool_name="my_tool",
+        input_policy="blocked",
+        output_policy="untrusted",
+        updated_by="admin",
+    )
+    prisma = _make_prisma(find_unique_row=row)
+    result = await update_tool_policy(
+        prisma, "my_tool", updated_by="admin", input_policy="blocked"
+    )
     assert result is not None
-    assert result.call_policy == "blocked"
-    prisma.db.execute_raw.assert_awaited_once()
-    call_args = prisma.db.execute_raw.call_args
-    sql = call_args.args[0]
-    assert "ON CONFLICT" in sql
-    assert "call_policy" in sql
+    assert result.input_policy == "blocked"
+    prisma.db.litellm_tooltable.upsert.assert_awaited_once()
+    call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs
+    assert call_kw["where"] == {"tool_name": "my_tool"}
+    assert call_kw["data"]["update"]["input_policy"] == "blocked"
+    assert call_kw["data"]["update"]["updated_by"] == "admin"
+    prisma.db.litellm_tooltable.find_unique.assert_awaited_with(
+        where={"tool_name": "my_tool"}
+    )
 
 
 @pytest.mark.asyncio
 async def test_get_tools_by_names_returns_policy_map():
     rows = [
-        {"tool_name": "tool_a", "call_policy": "trusted"},
-        {"tool_name": "tool_b", "call_policy": "blocked"},
+        _mock_row(tool_name="tool_a", input_policy="trusted", output_policy="untrusted"),
+        _mock_row(tool_name="tool_b", input_policy="blocked", output_policy="untrusted"),
     ]
-    prisma = _make_prisma(query_rows=rows)
+    prisma = _make_prisma(find_many_rows=rows)
     result = await get_tools_by_names(prisma, ["tool_a", "tool_b"])
-    assert result == {"tool_a": "trusted", "tool_b": "blocked"}
+    assert result == {
+        "tool_a": ("trusted", "untrusted"),
+        "tool_b": ("blocked", "untrusted"),
+    }
+    prisma.db.litellm_tooltable.find_many.assert_awaited_once_with(
+        where={"tool_name": {"in": ["tool_a", "tool_b"]}}
+    )
 
 
 @pytest.mark.asyncio
@@ -194,4 +221,71 @@ async def test_get_tools_by_names_empty_list():
     prisma = _make_prisma()
     result = await get_tools_by_names(prisma, [])
     assert result == {}
-    prisma.db.query_raw.assert_not_awaited()
+    prisma.db.litellm_tooltable.find_many.assert_not_awaited()
+
+
+# --- ToolPolicyRegistry ---
+
+
+def _mock_tool_row(
+    tool_name: str,
+    input_policy: str = "untrusted",
+    output_policy: str = "untrusted",
+):
+    row = MagicMock()
+    row.tool_name = tool_name
+    row.input_policy = input_policy
+    row.output_policy = output_policy
+    return row
+
+
+def _mock_perm_row(object_permission_id: str, blocked_tools: list):
+    row = MagicMock()
+    row.object_permission_id = object_permission_id
+    row.blocked_tools = blocked_tools
+    return row
+
+
+@pytest.mark.asyncio
+async def test_tool_policy_registry_sync_and_get_effective_policies():
+    """Registry syncs from DB; get_effective_policies returns merged blocked + global."""
+    prisma = MagicMock()
+    prisma.db.litellm_tooltable.find_many = AsyncMock(
+        return_value=[
+            _mock_tool_row("tool_a", input_policy="trusted"),
+            _mock_tool_row("tool_b", input_policy="blocked"),
+            _mock_tool_row("tool_c", input_policy="untrusted"),
+        ]
+    )
+    prisma.db.litellm_objectpermissiontable.find_many = AsyncMock(
+        return_value=[
+            _mock_perm_row("op-key-1", ["tool_a"]),
+            _mock_perm_row("op-team-1", ["tool_c"]),
+        ]
+    )
+    registry = get_tool_policy_registry()
+    await registry.sync_tool_policy_from_db(prisma)
+    assert registry.is_initialized()
+    # Key blocked: tool_a. Team blocked: tool_c. Global: tool_b blocked.
+    result = registry.get_effective_policies(
+        ["tool_a", "tool_b", "tool_c"],
+        object_permission_id="op-key-1",
+        team_object_permission_id="op-team-1",
+    )
+    assert result["tool_a"] == "blocked"
+    assert result["tool_b"] == "blocked"
+    assert result["tool_c"] == "blocked"
+    # No op ids: only global
+    result_global = registry.get_effective_policies(["tool_a", "tool_b", "tool_c"])
+    assert result_global["tool_a"] == "trusted"
+    assert result_global["tool_b"] == "blocked"
+    assert result_global["tool_c"] == "untrusted"
+
+
+@pytest.mark.asyncio
+async def test_tool_policy_registry_not_initialized_returns_untrusted():
+    """When not synced, get_effective_policies still returns untrusted for unknown tools."""
+    registry = ToolPolicyRegistry()
+    assert not registry.is_initialized()
+    result = registry.get_effective_policies(["unknown_tool"])
+    assert result == {"unknown_tool": "untrusted"}
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py
index c6a81efbf0b..943a8d4be75 100644
--- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py
@@ -12,9 +12,8 @@ from fastapi import HTTPException
 
 sys.path.insert(0, os.path.abspath("../../../../../.."))
 
-from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import (
-    ToolPolicyGuardrail,
-)
+from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import \
+    ToolPolicyGuardrail
 from litellm.types.guardrails import GuardrailEventHooks
 
 
@@ -70,10 +69,21 @@ async def test_no_tool_calls_in_response_passes_through(guardrail):
     assert result is inputs
 
 
+def _registry_mock(policy_map: dict):
+    """Return a mock registry with is_initialized=True and get_effective_policies returning policy_map."""
+    reg = MagicMock()
+    reg.is_initialized.return_value = True
+    reg.get_effective_policies.return_value = policy_map
+    return reg
+
+
 @pytest.mark.asyncio
 async def test_untrusted_tools_pass_through(guardrail):
     policy_map = {"search": "untrusted", "read_file": "trusted"}
-    with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
+    with patch(
+        "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
+        return_value=_registry_mock(policy_map),
+    ):
         inputs: Any = _tool_request_inputs(["search", "read_file"])
         result = await guardrail.apply_guardrail(
             inputs=inputs, request_data={}, input_type="request"
@@ -84,7 +94,10 @@ async def test_untrusted_tools_pass_through(guardrail):
 @pytest.mark.asyncio
 async def test_blocked_tool_in_request_raises_http_exception(guardrail):
     policy_map = {"dangerous_tool": "blocked"}
-    with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
+    with patch(
+        "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
+        return_value=_registry_mock(policy_map),
+    ):
         inputs: Any = _tool_request_inputs(["dangerous_tool"])
         with pytest.raises(HTTPException) as exc_info:
             await guardrail.apply_guardrail(
@@ -97,7 +110,10 @@ async def test_blocked_tool_in_request_raises_http_exception(guardrail):
 @pytest.mark.asyncio
 async def test_blocked_tool_in_response_raises_http_exception(guardrail):
     policy_map = {"exfil_tool": "blocked"}
-    with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
+    with patch(
+        "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
+        return_value=_registry_mock(policy_map),
+    ):
         inputs: Any = _tool_response_inputs(["exfil_tool"])
         with pytest.raises(HTTPException) as exc_info:
             await guardrail.apply_guardrail(
@@ -110,7 +126,10 @@ async def test_blocked_tool_in_response_raises_http_exception(guardrail):
 @pytest.mark.asyncio
 async def test_mixed_blocked_and_allowed_raises_for_blocked(guardrail):
     policy_map = {"safe_tool": "trusted", "bad_tool": "blocked"}
-    with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
+    with patch(
+        "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
+        return_value=_registry_mock(policy_map),
+    ):
         inputs: Any = _tool_request_inputs(["safe_tool", "bad_tool"])
         with pytest.raises(HTTPException) as exc_info:
             await guardrail.apply_guardrail(
@@ -123,8 +142,11 @@ async def test_mixed_blocked_and_allowed_raises_for_blocked(guardrail):
 
 @pytest.mark.asyncio
 async def test_tool_not_in_db_passes_through(guardrail):
-    """Tools not found in the DB (no entry) should not be blocked."""
-    with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value={})):
+    """When registry returns no policy (or empty), tools are not blocked."""
+    with patch(
+        "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
+        return_value=_registry_mock({}),
+    ):
         inputs: Any = _tool_request_inputs(["unknown_tool"])
         result = await guardrail.apply_guardrail(
             inputs=inputs, request_data={}, input_type="request"
@@ -133,43 +155,30 @@ async def test_tool_not_in_db_passes_through(guardrail):
 
 
 @pytest.mark.asyncio
-async def test_get_policies_cached_uses_cache(guardrail):
-    """Second call with same tool names should return the cached result."""
-    policy_map = {"tool_a": "trusted"}
+async def test_registry_not_initialized_passes_through(guardrail):
+    """When registry is not initialized, no tools are blocked (empty policy map)."""
+    reg = MagicMock()
+    reg.is_initialized.return_value = False
     with patch(
-        "litellm.proxy.db.tool_registry_writer.get_tools_by_names",
-        new=AsyncMock(return_value=policy_map),
-    ) as mock_db, patch(
-        "litellm.proxy.proxy_server.prisma_client",
-        new=MagicMock(),
+        "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
+        return_value=reg,
     ):
-        # first call — should hit DB
-        result1 = await guardrail._get_policies_cached(["tool_a"])
-        assert result1 == policy_map
-
-        # second call — should hit cache, not DB again
-        result2 = await guardrail._get_policies_cached(["tool_a"])
-        assert result2 == policy_map
-
-    assert mock_db.call_count == 1
-
-
-@pytest.mark.asyncio
-async def test_get_policies_cached_no_prisma(guardrail):
-    """Without a prisma client, returns empty dict."""
-    with patch(
-        "litellm.proxy.proxy_server.prisma_client",
-        None,
-    ):
-        result = await guardrail._get_policies_cached(["tool_a"])
-    assert result == {}
+        inputs: Any = _tool_request_inputs(["any_tool"])
+        result = await guardrail.apply_guardrail(
+            inputs=inputs, request_data={}, input_type="request"
+        )
+    assert result is inputs
+    reg.get_effective_policies.assert_not_called()
 
 
 @pytest.mark.asyncio
 async def test_response_tool_calls_as_objects(guardrail):
     """tool_calls that are objects (not dicts) with .function.name should work."""
     policy_map = {"obj_tool": "blocked"}
-    with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
+    with patch(
+        "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
+        return_value=_registry_mock(policy_map),
+    ):
         fn = MagicMock()
         fn.name = "obj_tool"
         tc = MagicMock()
diff --git a/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py b/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py
new file mode 100644
index 00000000000..4adc5acde8b
--- /dev/null
+++ b/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py
@@ -0,0 +1,200 @@
+"""
+Tests for tool allowlist enforcement (key/team metadata.allowed_tools).
+
+Covers:
+- check_tools_allowlist: allowed, disallowed, no allowlist, non-tool routes
+- extract_request_tool_names: OpenAI chat, responses, Anthropic, generate_content, MCP
+"""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from litellm.proxy._types import (ProxyErrorTypes, ProxyException,
+                                  UserAPIKeyAuth)
+from litellm.proxy.auth.auth_checks import check_tools_allowlist
+from litellm.proxy.guardrails.tool_name_extraction import (
+    TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names)
+
+
+def _token(metadata=None, team_metadata=None):
+    return UserAPIKeyAuth(
+        api_key="test-key",
+        user_id="user",
+        team_id="team",
+        org_id=None,
+        models=["*"],
+        metadata=metadata or {},
+        team_metadata=team_metadata or {},
+    )
+
+
+class TestExtractRequestToolNames:
+    """Test tool name extraction per API format."""
+
+    def test_openai_chat_tools(self):
+        data = {
+            "tools": [
+                {"type": "function", "function": {"name": "get_weather"}},
+                {"type": "function", "function": {"name": "run_sql"}},
+            ]
+        }
+        assert extract_request_tool_names("/v1/chat/completions", data) == [
+            "get_weather",
+            "run_sql",
+        ]
+
+    def test_openai_chat_functions_legacy(self):
+        data = {"functions": [{"name": "get_weather"}, {"name": "run_sql"}]}
+        assert extract_request_tool_names("/v1/chat/completions", data) == [
+            "get_weather",
+            "run_sql",
+        ]
+
+    def test_openai_responses_function_tools(self):
+        data = {
+            "tools": [
+                {"type": "function", "name": "get_current_weather", "description": "x"},
+            ]
+        }
+        assert extract_request_tool_names("/v1/responses", data) == [
+            "get_current_weather"
+        ]
+
+    def test_openai_responses_mcp_tools(self):
+        data = {
+            "tools": [
+                {"type": "mcp", "server_label": "dmcp", "server_url": "http://x"},
+            ]
+        }
+        assert extract_request_tool_names("/v1/responses", data) == ["dmcp"]
+
+    def test_anthropic_tools(self):
+        data = {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]}
+        assert extract_request_tool_names("/v1/messages", data) == [
+            "get_weather",
+            "run_sql",
+        ]
+
+    def test_generate_content_tools(self):
+        data = {
+            "tools": [
+                {
+                    "functionDeclarations": [
+                        {"name": "schedule_meeting", "description": "x"},
+                    ]
+                },
+            ]
+        }
+        assert extract_request_tool_names("/generate_content", data) == [
+            "schedule_meeting"
+        ]
+
+    def test_mcp_call_tool_name(self):
+        data = {"name": "my_tool", "arguments": {}}
+        assert extract_request_tool_names("/mcp/call_tool", data) == ["my_tool"]
+
+    def test_mcp_call_tool_mcp_tool_name(self):
+        data = {"mcp_tool_name": "other_tool"}
+        assert extract_request_tool_names("/mcp/call_tool", data) == ["other_tool"]
+
+    def test_non_tool_route_returns_empty(self):
+        data = {"tools": [{"type": "function", "function": {"name": "x"}}]}
+        assert extract_request_tool_names("/v1/embeddings", data) == []
+
+
+class TestCheckToolsAllowlist:
+    """Test allowlist enforcement in auth (no DB in hot path)."""
+
+    @pytest.mark.asyncio
+    async def test_no_allowlist_passes(self):
+        token = _token(metadata={}, team_metadata={})
+        body = {
+            "tools": [{"type": "function", "function": {"name": "get_weather"}}]
+        }
+        await check_tools_allowlist(
+            request_body=body,
+            valid_token=token,
+            team_object=None,
+            route="/v1/chat/completions",
+        )
+
+    @pytest.mark.asyncio
+    async def test_allowed_tool_passes(self):
+        token = _token(metadata={"allowed_tools": ["get_weather"]})
+        body = {
+            "tools": [{"type": "function", "function": {"name": "get_weather"}}]
+        }
+        await check_tools_allowlist(
+            request_body=body,
+            valid_token=token,
+            team_object=None,
+            route="/v1/chat/completions",
+        )
+
+    @pytest.mark.asyncio
+    async def test_disallowed_tool_raises(self):
+        token = _token(metadata={"allowed_tools": ["other_tool"]})
+        body = {
+            "tools": [{"type": "function", "function": {"name": "get_weather"}}]
+        }
+        with pytest.raises(ProxyException) as exc_info:
+            await check_tools_allowlist(
+                request_body=body,
+                valid_token=token,
+                team_object=None,
+                route="/v1/chat/completions",
+            )
+        assert exc_info.value.type == ProxyErrorTypes.tool_access_denied
+        assert "get_weather" in str(exc_info.value.message)
+
+    @pytest.mark.asyncio
+    async def test_team_allowlist_used_when_key_empty(self):
+        token = _token(
+            metadata={},
+            team_metadata={"allowed_tools": ["get_weather"]},
+        )
+        body = {
+            "tools": [{"type": "function", "function": {"name": "get_weather"}}]
+        }
+        await check_tools_allowlist(
+            request_body=body,
+            valid_token=token,
+            team_object=None,
+            route="/v1/chat/completions",
+        )
+
+    @pytest.mark.asyncio
+    async def test_key_allowlist_overrides_team(self):
+        token = _token(
+            metadata={"allowed_tools": ["get_weather"]},
+            team_metadata={"allowed_tools": ["other_tool"]},
+        )
+        body = {
+            "tools": [{"type": "function", "function": {"name": "get_weather"}}]
+        }
+        await check_tools_allowlist(
+            request_body=body,
+            valid_token=token,
+            team_object=None,
+            route="/v1/chat/completions",
+        )
+
+    @pytest.mark.asyncio
+    async def test_valid_token_none_skips(self):
+        await check_tools_allowlist(
+            request_body={"tools": [{"type": "function", "function": {"name": "x"}}]},
+            valid_token=None,
+            team_object=None,
+            route="/v1/chat/completions",
+        )
+
+    @pytest.mark.asyncio
+    async def test_no_tools_in_body_passes(self):
+        token = _token(metadata={"allowed_tools": ["get_weather"]})
+        await check_tools_allowlist(
+            request_body={"messages": []},
+            valid_token=token,
+            team_object=None,
+            route="/v1/chat/completions",
+        )
diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 622c3bf70a9..b927f312df8 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -39,7 +39,7 @@ import UserDashboard from "@/components/user_dashboard";
 import { AccessGroupsPage } from "@/components/AccessGroups/AccessGroupsPage";
 import { ProjectsPage } from "@/components/Projects/ProjectsPage";
 import VectorStoreManagement from "@/components/vector_store_management";
-import ToolPolicies from "@/components/ToolPolicies";
+import ToolPoliciesView from "@/components/ToolPoliciesView";
 import SpendLogsTable from "@/components/view_logs";
 import ViewUserDashboard from "@/components/view_users";
 import { ThemeProvider } from "@/contexts/ThemeContext";
@@ -549,7 +549,7 @@ function CreateKeyPageContent() {
                   ) : page == "vector-stores" ? (
                     
                   ) : page == "tool-policies" ? (
-                    
+                    
                   ) : page == "guardrails-monitor" ? (
                     
                   ) : page == "new_usage" ? (
diff --git a/ui/litellm-dashboard/src/components/ToolDetail.tsx b/ui/litellm-dashboard/src/components/ToolDetail.tsx
new file mode 100644
index 00000000000..ed0f866acb8
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/ToolDetail.tsx
@@ -0,0 +1,445 @@
+"use client";
+
+import { ArrowLeftOutlined, HistoryOutlined, ToolOutlined } from "@ant-design/icons";
+import { useQuery, useQueryClient } from "@tanstack/react-query";
+import { Button, Select, Spin } from "antd";
+import React, { useCallback, useEffect, useMemo, useState } from "react";
+import TeamDropdown from "@/components/common_components/team_dropdown";
+import { LogViewer } from "@/components/GuardrailsMonitor/LogViewer";
+import type { LogEntry } from "@/components/GuardrailsMonitor/mockData";
+import { PolicySelect } from "@/components/ToolPolicies/PolicySelect";
+import {
+  deleteToolPolicyOverride,
+  fetchToolDetail,
+  fetchToolPolicyOptions,
+  getToolUsageLogs,
+  keyListCall,
+  teamListCall,
+  updateToolPolicy,
+  type ToolPolicyOption,
+  type ToolPolicyOverrideRow,
+} from "@/components/networking";
+import type { Team } from "@/components/key_team_helpers/key_list";
+
+interface ToolDetailProps {
+  toolName: string;
+  onBack: () => void;
+  accessToken: string | null;
+}
+
+interface KeyOption {
+  token: string;
+  key_alias?: string;
+}
+
+const TOOL_DETAIL_QUERY_KEY = "tool-detail";
+
+const LOGS_PAGE_SIZE = 50;
+
+function getDefaultLogsDateRange(): { start: string; end: string } {
+  const end = new Date();
+  const start = new Date();
+  start.setDate(start.getDate() - 90);
+  const fmt = (d: Date) =>
+    d.toISOString().slice(0, 19).replace("T", " ");
+  return { start: fmt(start), end: fmt(end) };
+}
+
+export function ToolDetail({ toolName, onBack, accessToken }: ToolDetailProps) {
+  const queryClient = useQueryClient();
+  const [overrideSaving, setOverrideSaving] = useState(false);
+  const [inputPolicySaving, setInputPolicySaving] = useState(false);
+  const [outputPolicySaving, setOutputPolicySaving] = useState(false);
+  const [blockScope, setBlockScope] = useState<"team" | "key">("team");
+  const [blockTeamId, setBlockTeamId] = useState(null);
+  const [blockKey, setBlockKey] = useState(null);
+
+  const logsDateRange = useMemo(() => getDefaultLogsDateRange(), []);
+
+  const { data: detail, isLoading: detailLoading, error: detailError } = useQuery({
+    queryKey: [TOOL_DETAIL_QUERY_KEY, toolName],
+    queryFn: () => fetchToolDetail(accessToken!, toolName),
+    enabled: !!accessToken && !!toolName,
+  });
+
+  const { data: policyOptions } = useQuery({
+    queryKey: ["tool-policy-options"],
+    queryFn: () => fetchToolPolicyOptions(accessToken!),
+    enabled: !!accessToken,
+    staleTime: 60_000,
+  });
+
+  const { data: teamsData } = useQuery({
+    queryKey: ["teams-list-tool-detail"],
+    queryFn: () => teamListCall(accessToken!, null, null),
+    enabled: !!accessToken,
+  });
+
+  const { data: keysData } = useQuery({
+    queryKey: ["keys-list-tool-detail"],
+    queryFn: () => keyListCall(accessToken!, null, null, null, null, null, 1, 100),
+    enabled: !!accessToken,
+  });
+
+  const { data: logsData, isLoading: logsLoading } = useQuery({
+    queryKey: ["tool-usage-logs", toolName, logsDateRange.start, logsDateRange.end],
+    queryFn: () =>
+      getToolUsageLogs(accessToken!, toolName, {
+        page: 1,
+        pageSize: LOGS_PAGE_SIZE,
+        startDate: logsDateRange.start,
+        endDate: logsDateRange.end,
+      }),
+    enabled: !!accessToken && !!toolName,
+  });
+
+  const logs: LogEntry[] = useMemo(() => {
+    const list = logsData?.logs ?? [];
+    return list.map((l) => ({
+      id: l.id,
+      timestamp: l.timestamp,
+      action: "passed" as const,
+      model: l.model ?? undefined,
+      input_snippet: l.input_snippet ?? undefined,
+    }));
+  }, [logsData?.logs]);
+
+  const teams: Team[] = useMemo(() => {
+    const arr = Array.isArray(teamsData) ? teamsData : teamsData?.data ?? [];
+    return arr.map((t: { team_id?: string; id?: string; team_alias?: string }) => ({
+      team_id: t.team_id ?? t.id ?? "",
+      team_alias: t.team_alias ?? t.team_id ?? "",
+      models: [],
+      max_budget: null,
+      budget_duration: null,
+      tpm_limit: null,
+      rpm_limit: null,
+      organization_id: "",
+      created_at: "",
+      keys: [],
+      members_with_roles: [],
+      spend: 0,
+    }));
+  }, [teamsData]);
+
+  const keys: KeyOption[] = useMemo(() => {
+    const keysRes = keysData?.keys ?? keysData?.data ?? [];
+    return keysRes.map((k: { token?: string; api_key?: string; key_hash?: string; key_alias?: string }) => ({
+      token: k.token ?? k.api_key ?? k.key_hash ?? "",
+      key_alias: k.key_alias ?? (k.token ?? k.api_key ?? k.key_hash)?.toString?.()?.substring?.(0, 8),
+    }));
+  }, [keysData]);
+
+  const invalidateDetail = useCallback(() => {
+    queryClient.invalidateQueries({ queryKey: [TOOL_DETAIL_QUERY_KEY, toolName] });
+  }, [queryClient, toolName]);
+
+  const handleInputPolicyChange = useCallback(
+    async (_name: string, newPolicy: string) => {
+      if (!accessToken) return;
+      setInputPolicySaving(true);
+      try {
+        await updateToolPolicy(accessToken, toolName, { input_policy: newPolicy });
+        invalidateDetail();
+      } catch (e: unknown) {
+        alert(`Failed to update input policy: ${e instanceof Error ? e.message : String(e)}`);
+      } finally {
+        setInputPolicySaving(false);
+      }
+    },
+    [accessToken, toolName, invalidateDetail]
+  );
+
+  const handleOutputPolicyChange = useCallback(
+    async (_name: string, newPolicy: string) => {
+      if (!accessToken) return;
+      setOutputPolicySaving(true);
+      try {
+        await updateToolPolicy(accessToken, toolName, { output_policy: newPolicy });
+        invalidateDetail();
+      } catch (e: unknown) {
+        alert(`Failed to update output policy: ${e instanceof Error ? e.message : String(e)}`);
+      } finally {
+        setOutputPolicySaving(false);
+      }
+    },
+    [accessToken, toolName, invalidateDetail]
+  );
+
+  const handleAddOverride = useCallback(async () => {
+    if (!accessToken || !toolName) return;
+    const isTeam = blockScope === "team";
+    if (isTeam && !blockTeamId) return;
+    if (!isTeam && !blockKey?.token) return;
+    setOverrideSaving(true);
+    try {
+      await updateToolPolicy(accessToken, toolName, { input_policy: "blocked" }, {
+        team_id: isTeam ? blockTeamId : undefined,
+        key_hash: !isTeam ? blockKey!.token : undefined,
+        key_alias: !isTeam ? blockKey!.key_alias : undefined,
+      });
+      invalidateDetail();
+      setBlockTeamId(null);
+      setBlockKey(null);
+    } catch (e: unknown) {
+      alert(`Failed to add override: ${e instanceof Error ? e.message : String(e)}`);
+    } finally {
+      setOverrideSaving(false);
+    }
+  }, [accessToken, toolName, blockScope, blockTeamId, blockKey, invalidateDetail]);
+
+  const handleRemoveOverride = useCallback(
+    async (override: ToolPolicyOverrideRow) => {
+      if (!accessToken || !toolName) return;
+      setOverrideSaving(true);
+      try {
+        await deleteToolPolicyOverride(accessToken, toolName, {
+          team_id: override.team_id ?? undefined,
+          key_hash: override.key_hash ?? undefined,
+        });
+        invalidateDetail();
+      } catch (e: unknown) {
+        alert(`Failed to remove override: ${e instanceof Error ? e.message : String(e)}`);
+      } finally {
+        setOverrideSaving(false);
+      }
+    },
+    [accessToken, toolName, invalidateDetail]
+  );
+
+  if (detailLoading && !detail) {
+    return (
+      

+        
+      

+    );
+  }
+
+  if (detailError && !detail) {
+    return (
+      

+        
+        
Failed to load tool details.

+      

+    );
+  }
+
+  if (!detail) {
+    return null;
+  }
+
+  const { tool, overrides } = detail;
+
+  const inputDesc = policyOptions?.input_policies?.find(
+    (p) => p.value === tool.input_policy
+  )?.description;
+  const outputDesc = policyOptions?.output_policies?.find(
+    (p) => p.value === tool.output_policy
+  )?.description;
+
+  return (
+    

+      

+        
+
+        

+          

+            

+              
+              
{tool.tool_name}

+              
+                {tool.origin ?? "—"}
+              
+              
+                {(tool.call_count ?? 0).toLocaleString()} calls
+              
+            

+            

+              {tool.user_agent && (
+                

+                  
User Agent:

+                  
{tool.user_agent}

+                

+              )}
+              {tool.created_at && (
+                

+                  
First Discovered:

+                  
{new Date(tool.created_at).toLocaleString()}

+                

+              )}
+              {tool.last_used_at && (
+                

+                  
Last Used:

+                  
{new Date(tool.last_used_at).toLocaleString()}

+                

+              )}
+            

+          

+        

+      

+
+      

+        {/* Two-panel policy layout */}
+        

+          

+            
Input Policy

+            

+              {inputDesc ?? "Controls what data this tool is allowed to accept."}
+            

+            
+          

+
+          

+            
Output Policy

+            

+              {outputDesc ?? "Controls how this tool's output is trusted by downstream tools."}
+            

+            
+          

+        

+
+        {overrides.length > 0 && (
+          

+            
Blocked for team or key

+            
    
+              {overrides.map((ov) => (
+                
  • 
+                  
+                    {ov.team_id ? `Team: ${ov.team_id}` : ""}
+                    {ov.team_id && ov.key_hash ? " · " : ""}
+                    {ov.key_hash ? `Key: ${ov.key_alias || ov.key_hash.substring(0, 8)}` : ""}
+                    {!ov.team_id && !ov.key_hash ? "—" : ""}
+                  
+                  
+                
    • 
+              ))}
+            

+          

+        )}
+
+        

+          
Block for team or key

+          

+            

+              Scope
+              

+                
+                
+              

+            

+            

+              
+                {blockScope === "team" ? "Team" : "Key"}
+              
+              {blockScope === "team" ? (
+                 setBlockTeamId(id || null)}
+                />
+              ) : (
+                 onChange(toolName, v)}
-      onClick={(e) => e.stopPropagation()}
-      style={{
-        minWidth: 110,
-        fontWeight: 500,
-      }}
-      popupMatchSelectWidth={false}
-      options={POLICY_OPTIONS.map((o) => ({
-        value: o.value,
-        label: (
-          
-            
-            {o.label}
-          
-        ),
-      }))}
-    />
-  );
-};
-
-export const ToolPolicies: React.FC = ({ accessToken }) => {
+export const ToolPolicies: React.FC = ({ accessToken, onSelectTool }) => {
   const [tools, setTools] = useState([]);
   const [loading, setLoading] = useState(true);
   const [isFetching, setIsFetching] = useState(false);
   const [error, setError] = useState(null);
-  const [saving, setSaving] = useState(null);
+  const [savingInput, setSavingInput] = useState(null);
+  const [savingOutput, setSavingOutput] = useState(null);
 
   const [searchTerm, setSearchTerm] = useState("");
   const [sortField, setSortField] = useState("created_at");
@@ -123,16 +96,29 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
     return () => clearInterval(id);
   }, [isLiveTail, load]);
 
-  const handlePolicyChange = async (toolName: string, newPolicy: string) => {
+  const handleInputPolicyChange = async (toolName: string, newPolicy: string) => {
     if (!accessToken) return;
-    setSaving(toolName);
+    setSavingInput(toolName);
     try {
-      await updateToolPolicy(accessToken, toolName, newPolicy);
-      setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t)));
+      await updateToolPolicy(accessToken, toolName, { input_policy: newPolicy });
+      setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, input_policy: newPolicy } : t)));
     } catch (e: any) {
-      alert(`Failed to update policy: ${e.message}`);
+      alert(`Failed to update input policy: ${e.message}`);
     } finally {
-      setSaving(null);
+      setSavingInput(null);
+    }
+  };
+
+  const handleOutputPolicyChange = async (toolName: string, newPolicy: string) => {
+    if (!accessToken) return;
+    setSavingOutput(toolName);
+    try {
+      await updateToolPolicy(accessToken, toolName, { output_policy: newPolicy });
+      setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, output_policy: newPolicy } : t)));
+    } catch (e: any) {
+      alert(`Failed to update output policy: ${e.message}`);
+    } finally {
+      setSavingOutput(null);
     }
   };
 
@@ -157,7 +143,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
     setCurrentPage(1);
   };
 
-  // Build unique team/key options from loaded data
   const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map((v) => ({
     label: v as string,
     value: v as string,
@@ -169,9 +154,14 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
 
   const filterOptions: FilterOption[] = [
     {
-      name: "Policy",
-      label: "Policy",
-      options: POLICY_OPTIONS.map((o) => ({ label: o.label, value: o.value })),
+      name: "Input Policy",
+      label: "Input Policy",
+      options: INPUT_POLICY_OPTIONS.map((o) => ({ label: o.label, value: o.value })),
+    },
+    {
+      name: "Output Policy",
+      label: "Output Policy",
+      options: OUTPUT_POLICY_OPTIONS.map((o) => ({ label: o.label, value: o.value })),
     },
     {
       name: "Team Name",
@@ -185,6 +175,39 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
     },
   ];
 
+  const { newToday, newYesterday, trendSubtitle, totalTools, blockedCount, activeTeamsCount, needsReviewTools } =
+    useMemo(() => {
+      const now = new Date();
+      const todayKey = getUTCDateKey(now);
+      const yesterday = new Date(now);
+      yesterday.setUTCDate(yesterday.getUTCDate() - 1);
+      const yesterdayKey = getUTCDateKey(yesterday);
+
+      const newToday = countToolsInUTCDay(tools, todayKey);
+      const newYesterday = countToolsInUTCDay(tools, yesterdayKey);
+      const trendSubtitle = getTrendSubtitle(newToday, newYesterday);
+
+      const totalTools = tools.length;
+      const blockedCount = tools.filter((t) => t.input_policy === "blocked").length;
+      const activeTeamsCount = new Set(tools.map((t) => t.team_id).filter(Boolean)).size;
+
+      const needsReviewTools = tools.filter(
+        (t) =>
+          isCreatedInUTCDay(t.created_at, todayKey) &&
+          t.input_policy === "untrusted"
+      );
+
+      return {
+        newToday,
+        newYesterday,
+        trendSubtitle,
+        totalTools,
+        blockedCount,
+        activeTeamsCount,
+        needsReviewTools,
+      };
+    }, [tools]);
+
   const SortHeader = ({ label, field }: { label: string; field: SortField }) => (
     

       {label}
@@ -203,10 +226,12 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
         (t.team_id ?? "").toLowerCase().includes(q) ||
         (t.key_alias ?? "").toLowerCase().includes(q) ||
         (t.key_hash ?? "").toLowerCase().includes(q) ||
-        t.call_policy.toLowerCase().includes(q);
+        t.input_policy.toLowerCase().includes(q) ||
+        t.output_policy.toLowerCase().includes(q);
       if (!matchesSearch) return false;
     }
-    if (activeFilters["Policy"] && t.call_policy !== activeFilters["Policy"]) return false;
+    if (activeFilters["Input Policy"] && t.input_policy !== activeFilters["Input Policy"]) return false;
+    if (activeFilters["Output Policy"] && t.output_policy !== activeFilters["Output Policy"]) return false;
     if (activeFilters["Team Name"] && t.team_id !== activeFilters["Team Name"]) return false;
     if (activeFilters["Key Name"] && t.key_alias !== activeFilters["Key Name"]) return false;
     return true;
@@ -223,11 +248,74 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
   const totalPages = Math.max(1, Math.ceil(sorted.length / pageSize));
   const paginated = sorted.slice((currentPage - 1) * pageSize, currentPage * pageSize);
 
+  const scrollToToolRow = (toolId: string) => {
+    const idx = sorted.findIndex((t) => t.tool_id === toolId);
+    if (idx >= 0) {
+      const page = Math.floor(idx / pageSize) + 1;
+      if (page !== currentPage) setCurrentPage(page);
+      requestAnimationFrame(() => {
+        setTimeout(() => {
+          document.getElementById(`tool-row-${toolId}`)?.scrollIntoView({ behavior: "smooth", block: "center" });
+        }, 100);
+      });
+    }
+  };
+
   return (
-    

+    

       
Tool Policies

+
+      

+        
+              
+            
+          }
+        />
+        
+         0 ? "text-red-600" : undefined}
+        />
+         0 ? activeTeamsCount : "—"} />
+      

+
+      {needsReviewTools.length > 0 && (
+        

+          
Needs Review

+          

+            {needsReviewTools.length} new tool{needsReviewTools.length !== 1 ? "s" : ""} discovered that require
+            policy decisions.
+          

+          

+            {needsReviewTools.map((t) => (
+              
+                
+                  {t.tool_name}
+                
+                
+              
+            ))}
+          

+        

+      )}
+
       

-        {/* Toolbar */}
         

           

             

@@ -311,7 +399,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
             

           

 
-          {/* Filter row */}
           

              = ({ accessToken }) => {
           

         

 
-        {/* Auto-refresh banner */}
         {isLiveTail && (
           

             Auto-refreshing every 15 seconds
@@ -336,7 +422,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
           
{error}

         )}
 
-        {/* Table */}
         
             
@@ -347,7 +432,10 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
                 
               
               
-                
+                
+              
+              
+                
               
               
                 
@@ -359,45 +447,61 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
               
                 
               
-              Origin
+              User Agent
             
           
             {loading ? (
               
-                
+                
                   Loading tools…
                 
               
             ) : paginated.length === 0 ? (
               
-                
+                
                   No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery.
                 
               
             ) : (
               paginated.map((tool) => (
-                
+                
                   
                     
                   
                   
-                    
-                      
-                        {tool.tool_name}
-                      
-                    
+                    
                   
                   
                     
                   
-                  
-                    {(tool.call_count ?? 0).toLocaleString()}
+                  
+                    
+                  
+                  
+                    

+                      {(tool.call_count ?? 0).toLocaleString()}
+                    

                   
                   
                     
@@ -417,8 +521,8 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
                     
                   
                   
-                    
-                      {tool.origin ?? "-"}
+                    
+                      {tool.user_agent ?? "-"}
                     
                   
                 
@@ -427,7 +531,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
           

           
           
         

 
-        {/* Bottom pagination (only when > 1 page) */}
         {totalPages > 1 && (
           

             
@@ -453,6 +556,7 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
           

         )}
       

+
     

   );
 };
diff --git a/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.tsx b/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.tsx
new file mode 100644
index 00000000000..1317351931e
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.tsx
@@ -0,0 +1,92 @@
+"use client";
+
+import React from "react";
+import { Select } from "antd";
+
+export const INPUT_POLICY_OPTIONS = [
+  { value: "untrusted", label: "untrusted", color: "#92400e", bg: "#fef3c7", border: "#fcd34d" },
+  { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" },
+  { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" },
+] as const;
+
+export const OUTPUT_POLICY_OPTIONS = [
+  { value: "untrusted", label: "untrusted", color: "#92400e", bg: "#fef3c7", border: "#fcd34d" },
+  { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" },
+] as const;
+
+export const POLICY_OPTIONS = INPUT_POLICY_OPTIONS;
+
+export const policyStyle = (p: string) =>
+  INPUT_POLICY_OPTIONS.find((o) => o.value === p) ?? INPUT_POLICY_OPTIONS[0];
+
+export interface PolicySelectProps {
+  value: string;
+  toolName: string;
+  saving: boolean;
+  onChange: (toolName: string, policy: string) => void;
+  policyType?: "input" | "output";
+  size?: "small" | "middle";
+  minWidth?: number;
+  stopPropagation?: boolean;
+}
+
+export const PolicySelect: React.FC = ({
+  value,
+  toolName,
+  saving,
+  onChange,
+  policyType = "input",
+  size = "small",
+  minWidth = 110,
+  stopPropagation = true,
+}) => {
+  const options = policyType === "output" ? OUTPUT_POLICY_OPTIONS : INPUT_POLICY_OPTIONS;
+  const style = policyStyle(value);
+  return (
+