From 26611dd4bd953b4180a55d219d6f92fbd3955ca4 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 29 Jan 2026 16:40:55 -0800 Subject: [PATCH 001/132] fix: dead code cleanup in MCP server error handler raise e makes the error logging and JSON error response below it unreachable --- litellm/proxy/_experimental/mcp_server/server.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 6d54c3871e5..545d5956bd2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1898,7 +1898,6 @@ if MCP_AVAILABLE: await session_manager.handle_request(scope, receive, send) except Exception as e: - raise e verbose_logger.exception(f"Error handling MCP request: {e}") # Instead of re-raising, try to send a graceful error response try: From f9e8f8712b1b4eb8b619a9ff18630d84bfdfcc6e Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Feb 2026 16:51:45 -0800 Subject: [PATCH 002/132] fix: add cache invalidation for _cached_get_model_group_info on deployment changes _cached_get_model_group_info uses @lru_cache but had no invalidation, causing stale model group info (TPM/RPM limits) after dynamic deployment changes. Add cache_clear() at all 5 model_list mutation sites. --- litellm/router.py | 12 ++++++ tests/test_litellm/test_router.py | 67 +++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index d01c8443dab..5b72c3fb669 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6055,6 +6055,7 @@ class Router: self.model_list = [] self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index + self._invalidate_model_group_info_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works for model in original_model_list: @@ -6358,6 +6359,7 @@ class Router: """ idx = len(self.model_list) self.model_list.append(model) + self._invalidate_model_group_info_cache() # Update model_id index for O(1) lookup if model_id is not None: @@ -6405,6 +6407,7 @@ class Router: if removal_idx is not None: self.model_list.pop(removal_idx) + self._invalidate_model_group_info_cache() self._update_deployment_indices_after_removal( model_id=deployment_id, removal_idx=removal_idx ) @@ -6438,6 +6441,7 @@ class Router: if deployment_idx is not None: # Pop the item from the list first item = self.model_list.pop(deployment_idx) + self._invalidate_model_group_info_cache() self._update_deployment_indices_after_removal( model_id=id, removal_idx=deployment_idx ) @@ -7172,6 +7176,7 @@ class Router: """ # First populate the model_list self.model_list = [] + self._invalidate_model_group_info_cache() for _, model in enumerate(model_list): # Extract model_info from the model dict model_info = model.get("model_info", {}) @@ -7508,6 +7513,13 @@ class Router: return returned_models + def _invalidate_model_group_info_cache(self) -> None: + """Invalidate the cached model group info. + + Call this whenever self.model_list is modified to ensure the cache is rebuilt. + """ + self._cached_get_model_group_info.cache_clear() + def get_model_access_groups( self, model_name: Optional[str] = None, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 08ae804ea80..48aa435c3a9 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -925,6 +925,73 @@ def test_router_get_model_access_groups_team_only_models(): assert list(access_groups.keys()) == ["default-models"] +def test_cached_get_model_group_info(): + """ + Test that _cached_get_model_group_info caches results and + invalidates on deployment changes. + """ + from litellm.types.router import Deployment, LiteLLM_Params + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"tpm": 1000, "rpm": 100}, + }, + ] + ) + + # First call should compute and cache + result1 = router._cached_get_model_group_info("gpt-4") + assert result1 is not None + assert result1.tpm == 1000 + + # Second call should hit cache (same object) + result2 = router._cached_get_model_group_info("gpt-4") + assert result1 is result2 + + # Add a deployment — cache should be invalidated + router.add_deployment( + Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params(model="gpt-4", api_key="fake2"), + model_info={"tpm": 2000, "rpm": 200}, + ) + ) + result3 = router._cached_get_model_group_info("gpt-4") + assert result3 is not result2 + assert result3 is not None + assert result3.tpm == 3000 # 1000 + 2000 + + # Delete a deployment — cache should be invalidated + deployment_id = router.model_list[-1]["model_info"]["id"] + router.delete_deployment(id=deployment_id) + result4 = router._cached_get_model_group_info("gpt-4") + assert result4 is not result3 + assert result4 is not None + assert result4.tpm == 1000 + + # set_model_list — cache should be invalidated + router.set_model_list( + [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"tpm": 5000}, + }, + ] + ) + result5 = router._cached_get_model_group_info("gpt-4") + assert result5 is not result4 + assert result5 is not None + assert result5.tpm == 5000 + + # Verify cache still works after invalidation + result6 = router._cached_get_model_group_info("gpt-4") + assert result5 is result6 + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator(): """Test _acompletion_streaming_iterator for normal streaming and fallback behavior.""" From 6743d20de262338750f9b9046606222812915877 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 5 Feb 2026 10:50:45 -0800 Subject: [PATCH 003/132] docs: add trailing slash to /mcp endpoint URLs The /mcp endpoint requires a trailing slash because the MCP server is mounted as a sub-application using app.mount(). Starlette's mount behavior causes a 307 redirect from /mcp to /mcp/, which many MCP clients fail to handle. Updates documentation examples to use /mcp/ consistently. --- README.md | 2 +- docs/my-website/docs/mcp.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 77adddf8978..e7701b5cf9d 100644 --- a/README.md +++ b/README.md @@ -203,7 +203,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ { "mcpServers": { "LiteLLM": { - "url": "http://localhost:4000/mcp", + "url": "http://localhost:4000/mcp/", "headers": { "x-litellm-api-key": "Bearer sk-1234" } diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index d63b55ee29e..564a054aae2 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -632,7 +632,7 @@ import asyncio config = { "mcpServers": { "mcp_group": { - "url": "http://localhost:4000/mcp", + "url": "http://localhost:4000/mcp/", "headers": { "x-mcp-servers": "dev_group", # assume this gives access to github, zapier and deepwiki "x-litellm-api-key": "Bearer sk-1234", From 82f6d0fe431b1edb65ff6cd1c322fa7c588ba992 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Thu, 12 Feb 2026 14:47:09 -0800 Subject: [PATCH 004/132] healthcheck-model_id-fix: There was quite a bit of code that needed to be changed since health checks were entirely keyed by model name. This includes, proxy logic, the dashboard, and even networking, because model name was the identifier everywhere. All changed files had tests added to them, which are passing with no regressions. --- litellm/proxy/health_check.py | 12 +- .../health_endpoints/_health_endpoints.py | 8 +- .../litellm_utils_tests/test_health_check.py | 44 ++++ .../proxy/test_health_check_functions.py | 39 ++++ .../ModelsAndEndpointsView.test.tsx | 51 ++++- .../ModelsAndEndpointsView.tsx | 9 +- .../HealthCheckComponent.test.tsx | 147 +++++++++++++ .../model_dashboard/HealthCheckComponent.tsx | 199 ++++++++---------- .../model_dashboard/health_check_columns.tsx | 33 +-- .../src/components/networking.test.ts | 54 +++++ .../src/components/networking.tsx | 10 +- 11 files changed, 469 insertions(+), 137 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.test.tsx diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 427a16a9801..48a20833cb2 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -190,10 +190,15 @@ async def perform_health_check( model: Optional[str] = None, cli_model: Optional[str] = None, details: Optional[bool] = True, + model_id: Optional[str] = None, ): """ Perform a health check on the system. + When model_id is provided, only the deployment with that id is checked + (so models that share the same name but have different ids are checked separately). + When model (name) is provided, all deployments matching that name are checked. + Returns: (bool): True if the health check passes, False otherwise. """ @@ -205,7 +210,12 @@ async def perform_health_check( else: return [], [] - if model is not None: + # Filter by model_id first so a single deployment is checked when id is specified + if model_id is not None: + _by_id = [x for x in model_list if (x.get("model_info") or {}).get("id") == model_id] + if _by_id: + model_list = _by_id + elif model is not None: _new_model_list = [ x for x in model_list if x["litellm_params"]["model"] == model ] diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index da90696ec2d..b844a959622 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -238,7 +238,7 @@ async def health_services_endpoint( # noqa: PLR0915 service_in_success_callbacks = True else: for cb in litellm.success_callback: - if hasattr(cb, 'callback_name') and cb.callback_name == service: + if getattr(cb, "callback_name", None) == service: service_in_success_callbacks = True break cb_id = get_callback_identifier(cb) @@ -732,7 +732,11 @@ async def _perform_health_check_and_save( ): """Helper function to perform health check and save results to database""" healthy_endpoints, unhealthy_endpoints = await perform_health_check( - model_list=model_list, cli_model=cli_model, model=target_model, details=details + model_list=model_list, + cli_model=cli_model, + model=target_model, + details=details, + model_id=model_id, ) # Optionally save health check result to database (non-blocking) diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index 19882bbe4be..5123b0a7789 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -471,6 +471,50 @@ def test_update_litellm_params_for_health_check(): ) +@pytest.mark.asyncio +async def test_perform_health_check_filters_by_model_id(): + """ + When model_id is passed, only that deployment is checked (not all deployments + that share the same model name). + """ + from litellm.proxy.health_check import perform_health_check + + # Two deployments with same model_name but different ids + model_list = [ + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-id-1"}, + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, + }, + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-id-2"}, + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-2"}, + }, + ] + + captured_list = [] + + async def mock_perform_health_check(m_list, details=True): + captured_list.append(m_list) + return [{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}], [] + + with patch( + "litellm.proxy.health_check._perform_health_check", + side_effect=mock_perform_health_check, + ): + healthy_endpoints, unhealthy_endpoints = await perform_health_check( + model_list=model_list, model_id="deployment-id-2", details=True + ) + + # Only one deployment (deployment-id-2) should have been passed to _perform_health_check + assert len(captured_list) == 1 + assert len(captured_list[0]) == 1 + assert (captured_list[0][0].get("model_info") or {}).get("id") == "deployment-id-2" + assert len(healthy_endpoints) == 1 + assert healthy_endpoints[0]["api_key"] == "fake-key-2" + + @pytest.mark.asyncio async def test_perform_health_check_with_health_check_model(): """ diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index ccae9fb5425..4c91d0ae91e 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -12,6 +12,7 @@ sys.path.insert(0, os.path.abspath("../../..")) from litellm.proxy.health_endpoints._health_endpoints import ( _aggregate_health_check_results, _build_model_param_to_info_mapping, + _perform_health_check_and_save, _save_background_health_checks_to_db, _save_health_check_results_if_changed, _save_health_check_to_db, @@ -466,5 +467,43 @@ async def test_get_all_latest_health_checks_without_model_id(mock_prisma): assert result[0].checked_at == mock_check2.checked_at # Latest +@pytest.mark.asyncio +async def test_perform_health_check_and_save_passes_model_id_to_perform_health_check(): + """Test that _perform_health_check_and_save passes model_id to perform_health_check so health checks run by model id.""" + model_list = [ + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-abc"}, + "litellm_params": {"model": "gpt-4"}, + }, + ] + healthy = [{"model": "gpt-4"}] + unhealthy = [] + + async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None): + return healthy, unhealthy + + with patch( + "litellm.proxy.health_endpoints._health_endpoints.perform_health_check", + side_effect=mock_perform_health_check, + ) as mock_perform: + result = await _perform_health_check_and_save( + model_list=model_list, + target_model=None, + cli_model=None, + details=True, + prisma_client=None, + start_time=0.0, + user_id="user-1", + model_id="deployment-abc", + ) + + mock_perform.assert_called_once() + call_kwargs = mock_perform.call_args[1] + assert call_kwargs["model_id"] == "deployment-abc" + assert result["healthy_count"] == 1 + assert result["unhealthy_count"] == 0 + + if __name__ == "__main__": pytest.main([__file__]) \ No newline at end of file diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx index 1e8eabaea2e..b0df37ad6d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx @@ -1,6 +1,6 @@ /* @vitest-environment jsdom */ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render } from "@testing-library/react"; +import { act, render } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ModelsAndEndpointsView from "./ModelsAndEndpointsView"; @@ -13,6 +13,8 @@ vi.mock("@/components/networking", () => ({ getCallbacksCall: vi.fn().mockResolvedValue({ router_settings: {} }), setCallbacksCall: vi.fn().mockResolvedValue(undefined), getUiSettings: vi.fn().mockResolvedValue({ values: {} }), + latestHealthChecksCall: vi.fn().mockResolvedValue({ latest_health_checks: {} }), + getModelCostMapReloadStatus: vi.fn().mockResolvedValue({}), })); vi.mock("@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab", () => ({ @@ -27,6 +29,14 @@ vi.mock("@/components/add_model/AddModelForm", () => ({ default: () => null, })); +const mockHealthCheckComponent = vi.fn((_props: { all_models_on_proxy?: string[] }) => null); +vi.mock("@/components/model_dashboard/HealthCheckComponent", () => ({ + default: (props: { all_models_on_proxy?: string[] }) => { + mockHealthCheckComponent(props); + return null; + }, +})); + vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: () => ({ teams: [], @@ -104,4 +114,43 @@ describe("ModelsAndEndpointsView", () => { ); expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument(); }, 15000); + + it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => { + mockHealthCheckComponent.mockClear(); + const modelDataWithIds = { + data: [ + { model_name: "gpt-4", model_info: { id: "deployment-id-1" } }, + { model_name: "gpt-4", model_info: { id: "deployment-id-2" } }, + ], + }; + mockUseModelsInfo.mockReturnValue({ + data: { data: modelDataWithIds.data }, + isLoading: false, + refetch: vi.fn(), + }); + + const queryClient = createQueryClient(); + const { getByRole } = render( + + {}} + premiumUser={false} + teams={[]} + /> + , + ); + + const healthStatusTab = getByRole("tab", { name: "Health Status" }); + await act(async () => { + healthStatusTab.click(); + }); + + expect(mockHealthCheckComponent).toHaveBeenCalled(); + const healthCheckProps = mockHealthCheckComponent.mock.calls[0][0]; + expect(healthCheckProps.all_models_on_proxy).toEqual(["deployment-id-1", "deployment-id-2"]); + expect(healthCheckProps.all_models_on_proxy).not.toContain("gpt-4"); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 9d77774cb4c..b697a859dc5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -98,6 +98,13 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te return modelDataResponse.data.map((model: any) => model.model_name); }, [modelDataResponse?.data]); + const allModelIdsOnProxy = useMemo(() => { + if (!modelDataResponse?.data) return []; + return modelDataResponse.data + .map((model: any) => model.model_info?.id) + .filter((id: string | undefined): id is string => Boolean(id)); + }, [modelDataResponse?.data]); + const getProviderFromModel = (model: string) => { if (modelCostMapData !== null && modelCostMapData !== undefined) { if (typeof modelCostMapData == "object" && model in modelCostMapData) { @@ -397,7 +404,7 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te ({ + individualModelHealthCheckCall: (...args: unknown[]) => mockIndividualModelHealthCheckCall(...args), + latestHealthChecksCall: (...args: unknown[]) => mockLatestHealthChecksCall(...args), +})); + +describe("HealthCheckComponent", () => { + const getDisplayModelName = (model: { model_name?: string }) => model.model_name ?? ""; + + beforeEach(() => { + vi.clearAllMocks(); + mockLatestHealthChecksCall.mockResolvedValue({ latest_health_checks: {} }); + mockIndividualModelHealthCheckCall.mockResolvedValue({ + healthy_count: 1, + unhealthy_count: 0, + healthy_endpoints: [], + unhealthy_endpoints: [], + }); + }); + + it("should render the health check section", async () => { + const modelData = { + data: [ + { + model_name: "gpt-4", + model_info: { id: "deployment-1" }, + litellm_model_name: "gpt-4", + }, + ], + }; + + await act(async () => { + render( + , + ); + }); + await act(async () => { + await new Promise((r) => setTimeout(r, 0)); + }); + + expect(screen.getByText("Model Health Status")).toBeInTheDocument(); + expect( + screen.getByText("Run health checks on individual models to verify they are working correctly"), + ).toBeInTheDocument(); + }); + + it("should call individualModelHealthCheckCall with model id when run health check is triggered", async () => { + const modelData = { + data: [ + { + model_name: "gpt-4", + model_info: { id: "deployment-abc-123" }, + litellm_model_name: "gpt-4", + }, + ], + }; + + render( + , + ); + + const runButtons = screen.getAllByTestId("run-health-check-btn"); + expect(runButtons.length).toBeGreaterThanOrEqual(1); + const runButton = runButtons[0]; + + await act(async () => { + runButton.click(); + }); + + await act(async () => { + await new Promise((r) => setTimeout(r, 50)); + }); + + expect(mockIndividualModelHealthCheckCall).toHaveBeenCalledWith("token-123", "deployment-abc-123"); + expect(mockIndividualModelHealthCheckCall).not.toHaveBeenCalledWith("token-123", "gpt-4"); + }); + + it("should key health status by model id and show status from latest_health_checks by model_id", async () => { + const modelData = { + data: [ + { + model_name: "gpt-4", + model_info: { id: "id-alpha" }, + litellm_model_name: "gpt-4", + }, + { + model_name: "gpt-4", + model_info: { id: "id-beta" }, + litellm_model_name: "gpt-4", + }, + ], + }; + + mockLatestHealthChecksCall.mockResolvedValue({ + latest_health_checks: { + "id-alpha": { + status: "healthy", + checked_at: "2024-01-15T10:00:00Z", + error_message: null, + }, + "id-beta": { + status: "unhealthy", + checked_at: "2024-01-15T10:05:00Z", + error_message: "Connection failed", + }, + }, + }); + + await act(async () => { + render( + , + ); + }); + + await act(async () => { + await new Promise((r) => setTimeout(r, 0)); + }); + + expect(mockLatestHealthChecksCall).toHaveBeenCalledWith("token"); + const healthyBadges = screen.getAllByText("healthy"); + const unhealthyBadges = screen.getAllByText("unhealthy"); + expect(healthyBadges.length).toBeGreaterThanOrEqual(1); + expect(unhealthyBadges.length).toBeGreaterThanOrEqual(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx b/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx index 6fa83494408..b4bb1019dd6 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx @@ -53,31 +53,33 @@ const HealthCheckComponent: React.FC = ({ const healthTableRef = useRef>(null); - // Initialize health statuses on component mount + // Initialize health statuses on component mount (keyed by model id) useEffect(() => { if (!accessToken || !modelData?.data) return; const initializeHealthStatuses = async () => { const healthStatusMap: { [key: string]: HealthStatus } = {}; - // Initialize all models with default state using model names + // Initialize all models with default state using model ids modelData.data.forEach((model: any) => { - const modelName = model.model_name; - healthStatusMap[modelName] = { - status: "none", - lastCheck: "None", - lastSuccess: "None", - loading: false, - error: undefined, - fullError: undefined, - successResponse: undefined, - }; + const modelId = model.model_info?.id; + if (modelId) { + healthStatusMap[modelId] = { + status: "none", + lastCheck: "None", + lastSuccess: "None", + loading: false, + error: undefined, + fullError: undefined, + successResponse: undefined, + }; + } }); try { const latestHealthChecks = await latestHealthChecksCall(accessToken); - // Override with actual database data if it exists + // Override with actual database data if it exists (latest_health_checks is keyed by model_id) if ( latestHealthChecks && latestHealthChecks.latest_health_checks && @@ -86,32 +88,28 @@ const HealthCheckComponent: React.FC = ({ Object.entries(latestHealthChecks.latest_health_checks).forEach(([key, checkData]: [string, any]) => { if (!checkData) return; - let targetModelName: string | null = null; + let targetModelId: string | null = null; - // The key could be either model_id or model_name, try both approaches - const directModelMatch = modelData.data.find((m: any) => m.model_name === key); - if (directModelMatch) { - targetModelName = directModelMatch.model_name; + // The key is model_id from the backend; fallback to matching by model_name for legacy data + const modelByIdMatch = modelData.data.find((m: any) => m.model_info && m.model_info.id === key); + if (modelByIdMatch) { + targetModelId = modelByIdMatch.model_info.id; } else { - // If not a direct match, treat as model_id and find the corresponding model - const modelByIdMatch = modelData.data.find((m: any) => m.model_info && m.model_info.id === key); - if (modelByIdMatch) { - targetModelName = modelByIdMatch.model_name; - } else { - // Check if checkData contains model_name and use that - if (checkData.model_name) { - const modelByNameInData = modelData.data.find((m: any) => m.model_name === checkData.model_name); - if (modelByNameInData) { - targetModelName = modelByNameInData.model_name; - } + const directModelMatch = modelData.data.find((m: any) => m.model_name === key); + if (directModelMatch?.model_info?.id) { + targetModelId = directModelMatch.model_info.id; + } else if (checkData.model_name) { + const modelByNameInData = modelData.data.find((m: any) => m.model_name === checkData.model_name); + if (modelByNameInData?.model_info?.id) { + targetModelId = modelByNameInData.model_info.id; } } } - if (targetModelName) { + if (targetModelId) { const fullError = checkData.error_message || undefined; - healthStatusMap[targetModelName] = { + healthStatusMap[targetModelId] = { status: checkData.status || "unknown", lastCheck: checkData.checked_at ? new Date(checkData.checked_at).toLocaleString() : "None", lastSuccess: @@ -246,33 +244,31 @@ const HealthCheckComponent: React.FC = ({ return cleaned.length > 100 ? cleaned.substring(0, 97) + "..." : cleaned; }; - const runIndividualHealthCheck = async (modelName: string) => { + const runIndividualHealthCheck = async (modelId: string) => { if (!accessToken) return; setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { - ...prev[modelName], + [modelId]: { + ...prev[modelId], loading: true, status: "checking", }, })); try { - // Run the health check and process the response directly - const response = await individualModelHealthCheckCall(accessToken, modelName); + const response = await individualModelHealthCheckCall(accessToken, modelId); const currentTime = new Date().toLocaleString(); - // Check if there are any unhealthy endpoints (which means this specific model failed) if (response.unhealthy_count > 0 && response.unhealthy_endpoints && response.unhealthy_endpoints.length > 0) { const rawError = response.unhealthy_endpoints[0]?.error || "Health check failed"; const errorMessage = extractMeaningfulError(rawError); setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "unhealthy", lastCheck: currentTime, - lastSuccess: prev[modelName]?.lastSuccess || "None", + lastSuccess: prev[modelId]?.lastSuccess || "None", loading: false, error: errorMessage, fullError: rawError, @@ -281,7 +277,7 @@ const HealthCheckComponent: React.FC = ({ } else { setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "healthy", lastCheck: currentTime, lastSuccess: currentTime, @@ -291,41 +287,33 @@ const HealthCheckComponent: React.FC = ({ })); } - // Refresh health status from database to get the saved check data including timestamp try { const latestHealthChecks = await latestHealthChecksCall(accessToken); + const checkData = latestHealthChecks.latest_health_checks?.[modelId]; - // Find the model ID for this model name to look up database data - const model = modelData.data.find((m: any) => m.model_name === modelName); - if (model) { - const modelId = model.model_info.id; - const checkData = latestHealthChecks.latest_health_checks?.[modelId]; - - if (checkData) { - const fullError = checkData.error_message || undefined; - setModelHealthStatuses((prev) => ({ - ...prev, - [modelName]: { - status: checkData.status || prev[modelName]?.status || "unknown", - lastCheck: checkData.checked_at - ? new Date(checkData.checked_at).toLocaleString() - : prev[modelName]?.lastCheck || "None", - lastSuccess: - checkData.status === "healthy" - ? checkData.checked_at - ? new Date(checkData.checked_at).toLocaleString() - : prev[modelName]?.lastSuccess || "None" - : prev[modelName]?.lastSuccess || "None", - loading: false, - error: fullError ? extractMeaningfulError(fullError) : prev[modelName]?.error, - fullError: fullError || prev[modelName]?.fullError, - successResponse: checkData.status === "healthy" ? checkData : prev[modelName]?.successResponse, - }, - })); - } + if (checkData) { + const fullError = checkData.error_message || undefined; + setModelHealthStatuses((prev) => ({ + ...prev, + [modelId]: { + status: checkData.status || prev[modelId]?.status || "unknown", + lastCheck: checkData.checked_at + ? new Date(checkData.checked_at).toLocaleString() + : prev[modelId]?.lastCheck || "None", + lastSuccess: + checkData.status === "healthy" + ? checkData.checked_at + ? new Date(checkData.checked_at).toLocaleString() + : prev[modelId]?.lastSuccess || "None" + : prev[modelId]?.lastSuccess || "None", + loading: false, + error: fullError ? extractMeaningfulError(fullError) : prev[modelId]?.error, + fullError: fullError || prev[modelId]?.fullError, + successResponse: checkData.status === "healthy" ? checkData : prev[modelId]?.successResponse, + }, + })); } } catch (dbError) { - // Ignore database errors - we already have the health check result from the API call console.debug("Could not fetch updated status from database (non-critical):", dbError); } } catch (error) { @@ -334,10 +322,10 @@ const HealthCheckComponent: React.FC = ({ const errorMessage = extractMeaningfulError(rawError); setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "unhealthy", lastCheck: currentTime, - lastSuccess: prev[modelName]?.lastSuccess || "None", + lastSuccess: prev[modelId]?.lastSuccess || "None", loading: false, error: errorMessage, fullError: rawError, @@ -349,11 +337,10 @@ const HealthCheckComponent: React.FC = ({ const runAllHealthChecks = async () => { const modelsToCheck = selectedModelsForHealth.length > 0 ? selectedModelsForHealth : all_models_on_proxy; - // Set all models to loading state const loadingStatuses = modelsToCheck.reduce( - (acc, modelName) => { - acc[modelName] = { - ...modelHealthStatuses[modelName], + (acc, modelId) => { + acc[modelId] = { + ...modelHealthStatuses[modelId], loading: true, status: "checking", }; @@ -364,30 +351,25 @@ const HealthCheckComponent: React.FC = ({ setModelHealthStatuses((prev) => ({ ...prev, ...loadingStatuses })); - // Store results from individual health checks const healthCheckResults: { [key: string]: any } = {}; - // Run all health checks in parallel and collect results - const healthCheckPromises = modelsToCheck.map(async (modelName) => { + const healthCheckPromises = modelsToCheck.map(async (modelId) => { if (!accessToken) return; try { - // Run the health check and store the result - const response = await individualModelHealthCheckCall(accessToken, modelName); - healthCheckResults[modelName] = response; + const response = await individualModelHealthCheckCall(accessToken, modelId); + healthCheckResults[modelId] = response; - // Update status immediately based on response const currentTime = new Date().toLocaleString(); - // Check if there are any unhealthy endpoints (which means this specific model failed) if (response.unhealthy_count > 0 && response.unhealthy_endpoints && response.unhealthy_endpoints.length > 0) { const rawError = response.unhealthy_endpoints[0]?.error || "Health check failed"; const errorMessage = extractMeaningfulError(rawError); setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "unhealthy", lastCheck: currentTime, - lastSuccess: prev[modelName]?.lastSuccess || "None", + lastSuccess: prev[modelId]?.lastSuccess || "None", loading: false, error: errorMessage, fullError: rawError, @@ -396,7 +378,7 @@ const HealthCheckComponent: React.FC = ({ } else { setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "healthy", lastCheck: currentTime, lastSuccess: currentTime, @@ -406,17 +388,16 @@ const HealthCheckComponent: React.FC = ({ })); } } catch (error) { - console.error(`Health check failed for ${modelName}:`, error); - // Set error status for failed health checks + console.error(`Health check failed for model id ${modelId}:`, error); const currentTime = new Date().toLocaleString(); const rawError = error instanceof Error ? error.message : String(error); const errorMessage = extractMeaningfulError(rawError); setModelHealthStatuses((prev) => ({ ...prev, - [modelName]: { + [modelId]: { status: "unhealthy", lastCheck: currentTime, - lastSuccess: prev[modelName]?.lastSuccess || "None", + lastSuccess: prev[modelId]?.lastSuccess || "None", loading: false, error: errorMessage, fullError: rawError, @@ -425,27 +406,21 @@ const HealthCheckComponent: React.FC = ({ } }); - // Wait for all health checks to complete await Promise.allSettled(healthCheckPromises); - // Refresh health statuses from database to get the saved check data including timestamps try { if (!accessToken) return; const latestHealthChecks = await latestHealthChecksCall(accessToken); if (latestHealthChecks.latest_health_checks) { - // Update health statuses from database, which should have the most accurate saved data Object.entries(latestHealthChecks.latest_health_checks).forEach(([modelId, checkData]: [string, any]) => { - // Find the model name for this model ID - const model = modelData.data.find((m: any) => m.model_info.id === modelId); - if (model && modelsToCheck.includes(model.model_name) && checkData) { - const modelName = model.model_name; + if (modelsToCheck.includes(modelId) && checkData) { const fullError = checkData.error_message || undefined; setModelHealthStatuses((prev) => { - const currentStatus = prev[modelName]; + const currentStatus = prev[modelId]; return { ...prev, - [modelName]: { + [modelId]: { status: checkData.status || currentStatus?.status || "unknown", lastCheck: checkData.checked_at ? new Date(checkData.checked_at).toLocaleString() @@ -468,15 +443,14 @@ const HealthCheckComponent: React.FC = ({ } } catch (dbError) { console.warn("Failed to fetch updated health statuses from database (non-critical):", dbError); - // This is non-critical - we already have the health check results from the API calls } }; - const handleModelSelection = (modelName: string, checked: boolean) => { + const handleModelSelection = (modelId: string, checked: boolean) => { if (checked) { - setSelectedModelsForHealth((prev) => [...prev, modelName]); + setSelectedModelsForHealth((prev) => [...prev, modelId]); } else { - setSelectedModelsForHealth((prev) => prev.filter((name) => name !== modelName)); + setSelectedModelsForHealth((prev) => prev.filter((id) => id !== modelId)); setAllModelsSelected(false); } }; @@ -580,8 +554,9 @@ const HealthCheckComponent: React.FC = ({ teams, )} data={modelData.data.map((model: any) => { - const modelName = model.model_name; - const healthStatus = modelHealthStatuses[modelName] || { + const modelId = model.model_info?.id; + const healthStatus = modelId ? modelHealthStatuses[modelId] : null; + const status = healthStatus || { status: "none", lastCheck: "None", loading: false, @@ -591,12 +566,12 @@ const HealthCheckComponent: React.FC = ({ model_info: model.model_info, provider: model.provider, litellm_model_name: model.litellm_model_name, - health_status: healthStatus.status, - last_check: healthStatus.lastCheck, - last_success: healthStatus.lastSuccess || "None", - health_loading: healthStatus.loading, - health_error: healthStatus.error, - health_full_error: healthStatus.fullError, + health_status: status.status, + last_check: status.lastCheck, + last_success: status.lastSuccess || "None", + health_loading: status.loading, + health_error: status.error, + health_full_error: status.fullError, }; })} isLoading={false} diff --git a/ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx b/ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx index 3e8ae662ad4..396afb7ed0b 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx @@ -35,9 +35,9 @@ export const healthCheckColumns = ( modelHealthStatuses: { [key: string]: HealthStatus }, selectedModelsForHealth: string[], allModelsSelected: boolean, - handleModelSelection: (modelName: string, checked: boolean) => void, + handleModelSelection: (modelId: string, checked: boolean) => void, handleSelectAll: (checked: boolean) => void, - runIndividualHealthCheck: (modelName: string) => void, + runIndividualHealthCheck: (modelId: string) => void, getStatusBadge: (status: string) => JSX.Element, getDisplayModelName: (model: any) => string, showErrorModal?: (modelName: string, cleanedError: string, fullError: string) => void, @@ -62,14 +62,14 @@ export const healthCheckColumns = ( sortingFn: "alphanumeric", cell: ({ row }) => { const model = row.original; - const modelName = model.model_name; - const isSelected = selectedModelsForHealth.includes(modelName); + const modelId = model.model_info?.id ?? ""; + const isSelected = selectedModelsForHealth.includes(modelId); return (
handleModelSelection(modelName, e.target.checked)} + onChange={(e) => handleModelSelection(modelId, e.target.checked)} onClick={(e) => e.stopPropagation()} /> @@ -169,8 +169,9 @@ export const healthCheckColumns = ( ); } - const modelName = model.model_name; - const hasSuccessResponse = healthStatus.status === "healthy" && modelHealthStatuses[modelName]?.successResponse; + const modelId = model.model_info?.id ?? ""; + const displayName = getDisplayModelName(model) || model.model_name; + const hasSuccessResponse = healthStatus.status === "healthy" && modelHealthStatuses[modelId]?.successResponse; return (
@@ -178,7 +179,7 @@ export const healthCheckColumns = ( {hasSuccessResponse && showSuccessModal && ( )}
From ef67b6b53363ea064deba1914568044b54020efa Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 17:48:55 +0530 Subject: [PATCH 013/132] Add support for phase param --- litellm/types/responses/main.py | 3 + .../test_openai_responses_transformation.py | 316 +++++++++++++++++- 2 files changed, 318 insertions(+), 1 deletion(-) diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 8f6333ff900..bda53bae082 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -6,6 +6,7 @@ from typing_extensions import Any, List, Optional, TypedDict from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject +Phase = Optional[Literal["commentary", "final_answer"]] # TODO: Once openai sdk has updated, we can remove this and use the openai sdk type class GenericResponseOutputItemContentAnnotation(BaseLiteLLMOpenAIResponseObject): """Annotation for content in a message""" @@ -35,6 +36,7 @@ class OutputFunctionToolCall(BaseLiteLLMOpenAIResponseObject): type: Optional[str] # "function_call" id: Optional[str] status: Literal["in_progress", "completed", "incomplete"] + phase: Phase = None class OutputImageGenerationCall(BaseLiteLLMOpenAIResponseObject): @@ -57,6 +59,7 @@ class GenericResponseOutputItem(BaseLiteLLMOpenAIResponseObject): status: str # "completed", "in_progress", etc. role: str # "assistant", "user", etc. content: List[OutputText] + phase: Phase = None class DeleteResponseResult(BaseLiteLLMOpenAIResponseObject): diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index 7c08716c04c..1a5ab808f7b 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -925,4 +925,318 @@ def test_get_supported_openai_params(): assert "temperature" in params assert "stream" in params assert "background" in params - assert "stream" in params \ No newline at end of file + assert "stream" in params + + +class TestPhaseParameter: + """Tests for the `phase` parameter on assistant output items (gpt-5.3-codex).""" + + def setup_method(self): + self.config = OpenAIResponsesAPIConfig() + self.model = "gpt-5.3-codex" + self.logging_obj = MagicMock() + + @staticmethod + def _make_output_text(text: str): + from litellm.types.responses.main import OutputText + + return OutputText(type="output_text", text=text, annotations=[]) + + def test_generic_response_output_item_accepts_phase_commentary(self): + from litellm.types.responses.main import GenericResponseOutputItem + + item = GenericResponseOutputItem( + type="message", + id="msg_001", + status="completed", + role="assistant", + content=[self._make_output_text("Thinking...")], + phase="commentary", + ) + assert item.phase == "commentary" + + def test_generic_response_output_item_accepts_phase_final_answer(self): + from litellm.types.responses.main import GenericResponseOutputItem + + item = GenericResponseOutputItem( + type="message", + id="msg_002", + status="completed", + role="assistant", + content=[self._make_output_text("The answer is 42.")], + phase="final_answer", + ) + assert item.phase == "final_answer" + + def test_generic_response_output_item_phase_defaults_to_none(self): + from litellm.types.responses.main import GenericResponseOutputItem + + item = GenericResponseOutputItem( + type="message", + id="msg_003", + status="completed", + role="assistant", + content=[self._make_output_text("Hello")], + ) + assert item.phase is None + + def test_output_function_tool_call_accepts_phase(self): + from litellm.types.responses.main import OutputFunctionToolCall + + item = OutputFunctionToolCall( + type="function_call", + id="fc_001", + arguments='{"query": "test"}', + call_id="call_001", + name="search", + status="completed", + phase="commentary", + ) + assert item.phase == "commentary" + + def test_input_passthrough_dict_preserves_phase(self): + """Dict input items (the normal HTTP flow) must preserve phase verbatim.""" + input_items = [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Hi"}], + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Preamble..."}], + "phase": "commentary", + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done."}], + "phase": "final_answer", + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Neutral."}], + "phase": None, + }, + ] + + result = self.config._validate_input_param(input_items) + assert isinstance(result, list) + + assert "phase" not in result[0] + assert result[1]["phase"] == "commentary" + assert result[2]["phase"] == "final_answer" + assert result[3]["phase"] is None + + def test_input_passthrough_pydantic_preserves_non_null_phase(self): + """Pydantic input items must preserve non-null phase values.""" + from litellm.types.responses.main import GenericResponseOutputItem + + item = GenericResponseOutputItem( + type="message", + id="msg_010", + status="completed", + role="assistant", + content=[self._make_output_text("commentary")], + phase="commentary", + ) + + result = self.config._validate_input_param([item]) + assert isinstance(result, list) + assert result[0]["phase"] == "commentary" + + def test_response_parsing_preserves_phase_on_output(self): + """Non-streaming response must preserve phase on output items.""" + raw_json = { + "id": "resp_001", + "created_at": 1700000000, + "model": "gpt-5.3-codex", + "object": "response", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_001", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "preamble"}], + "phase": "commentary", + }, + { + "type": "message", + "id": "msg_002", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer"}], + "phase": "final_answer", + }, + ], + "usage": {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, + } + + response = ResponsesAPIResponse(**raw_json) + assert len(response.output) == 2 + + for idx, output_item in enumerate(response.output): + if isinstance(output_item, dict): + phase = output_item.get("phase") + else: + phase = getattr(output_item, "phase", None) + + expected = "commentary" if idx == 0 else "final_answer" + assert phase == expected, ( + f"output[{idx}] phase={phase!r}, expected {expected!r}" + ) + + def test_streaming_output_item_done_preserves_phase(self): + """OutputItemDoneEvent must preserve phase on its item.""" + from litellm.types.llms.openai import ( + OutputItemDoneEvent, + ResponsesAPIStreamEvents, + ) + + chunk = { + "type": "response.output_item.done", + "output_index": 0, + "sequence_number": 3, + "item": { + "type": "message", + "id": "msg_100", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "done"}], + "phase": "final_answer", + }, + } + + result = self.config.transform_streaming_response( + model=self.model, parsed_chunk=chunk, logging_obj=self.logging_obj + ) + + assert isinstance(result, OutputItemDoneEvent) + assert result.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE + assert getattr(result.item, "phase", None) == "final_answer" + + def test_streaming_output_item_added_preserves_phase(self): + """OutputItemAddedEvent must preserve phase on its item.""" + from litellm.types.llms.openai import ( + OutputItemAddedEvent, + ResponsesAPIStreamEvents, + ) + + chunk = { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "type": "message", + "id": "msg_200", + "role": "assistant", + "phase": "commentary", + }, + } + + result = self.config.transform_streaming_response( + model=self.model, parsed_chunk=chunk, logging_obj=self.logging_obj + ) + + assert isinstance(result, OutputItemAddedEvent) + assert result.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + assert getattr(result.item, "phase", None) == "commentary" + + def test_streaming_response_completed_preserves_phase(self): + """ResponseCompletedEvent must preserve phase on output items inside the response.""" + completed_chunk = { + "type": "response.completed", + "response": { + "id": "resp_300", + "created_at": 1700000000, + "model": "gpt-5.3-codex", + "object": "response", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_300", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "final"}], + "phase": "final_answer", + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 10, + "total_tokens": 15, + }, + }, + } + + result = self.config.transform_streaming_response( + model=self.model, + parsed_chunk=completed_chunk, + logging_obj=self.logging_obj, + ) + + assert result.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + output_item = result.response.output[0] + if isinstance(output_item, dict): + assert output_item["phase"] == "final_answer" + else: + assert getattr(output_item, "phase", None) == "final_answer" + + def test_phase_roundtrip_output_to_input(self): + """Simulate full round-trip: parse response output, then send items back as input.""" + raw_json = { + "id": "resp_rt", + "created_at": 1700000000, + "model": "gpt-5.3-codex", + "object": "response", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_rt1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "preamble"}], + "phase": "commentary", + }, + { + "type": "message", + "id": "msg_rt2", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer"}], + "phase": "final_answer", + }, + ], + "usage": {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, + } + + response = ResponsesAPIResponse(**raw_json) + + input_items = [] + for item in response.output: + if isinstance(item, dict): + input_items.append(item) + else: + input_items.append( + item.model_dump() if hasattr(item, "model_dump") else dict(item) + ) + + input_items.append( + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "next question"}], + } + ) + + validated = self.config._validate_input_param(input_items) + assert isinstance(validated, list) + + assert validated[0]["phase"] == "commentary" + assert validated[1]["phase"] == "final_answer" + assert "phase" not in validated[2] \ No newline at end of file From ac720defc3d7ae338299833055474dedfe667387 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 17:50:38 +0530 Subject: [PATCH 014/132] Add documentation related to phase --- docs/my-website/blog/gpt_5_3_codex/index.md | 145 ++++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 docs/my-website/blog/gpt_5_3_codex/index.md diff --git a/docs/my-website/blog/gpt_5_3_codex/index.md b/docs/my-website/blog/gpt_5_3_codex/index.md new file mode 100644 index 00000000000..321e573cccf --- /dev/null +++ b/docs/my-website/blog/gpt_5_3_codex/index.md @@ -0,0 +1,145 @@ +--- +slug: gpt_5_3_codex +title: "Day 0 Support: GPT-5.3-Codex" +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 +description: "Day 0 support for GPT-5.3-Codex on LiteLLM, including phase parameter handling for Responses API." +tags: [openai, gpt-5.3-codex, codex, day 0 support] +hide_table_of_contents: false +--- + +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +LiteLLM now supports GPT-5.3-Codex on Day 0, including support for the new assistant `phase` metadata on Responses API output items. + +## Why `phase` matters for GPT-5.3-Codex + +`phase` appears on assistant output items and helps distinguish preamble/commentary turns from final closeout responses. + +Reference: [Phase parameter docs](https://developers-site-git-alphas-venusaur-api-openai.vercel.app/alphas/venusaur-api/phase-parameter) + +Supported values: +- `null` +- `"commentary"` +- `"final_answer"` + +Important: +- Persist assistant output items with `phase` exactly as returned. +- Send those assistant items back on the next turn. +- Do **not** add `phase` to user messages. + +## Docker Image + +```bash +docker pull ghcr.io/berriai/litellm:v1.81.3-stable.sonnet-4-6 +``` + +## Usage + + + + +**1. Setup config.yaml** + +```yaml +model_list: + - model_name: gpt-5.3-codex + litellm_params: + model: openai/gpt-5.3-codex +``` + +**2. Start the proxy** + +```bash +docker run -d \ + -p 4000:4000 \ + -e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \ + -v $(pwd)/config.yaml:/app/config.yaml \ + ghcr.io/berriai/litellm:v1.81.3-stable.sonnet-4-6 \ + --config /app/config.yaml +``` + + +**3. Test it** + +```bash +curl -X POST "http://0.0.0.0:4000/v1/responses" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -d '{ + "model": "gpt-5.3-codex", + "input": "Write a Python script that checks if a number is prime." + }' +``` + + + + +## Python Example: Persist `phase` with OpenAI Client + LiteLLM Base URL + +```python +from openai import OpenAI + +client = OpenAI( + base_url="http://0.0.0.0:4000/v1", # LiteLLM Proxy + api_key="your-litellm-api-key", +) + +items = [] # Persist this per conversation/thread + + +def _item_get(item, key, default=None): + if isinstance(item, dict): + return item.get(key, default) + return getattr(item, key, default) + + +def run_turn(user_text: str): + global items + + # User message: no phase field + items.append( + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": user_text}], + } + ) + + resp = client.responses.create( + model="gpt-5.3-codex", + input=items, + ) + + # Persist assistant output items verbatim, including phase + for out_item in (resp.output or []): + items.append(out_item) + + # Optional: inspect latest phase for UI/telemetry routing + latest_phase = None + for out_item in reversed(resp.output or []): + if _item_get(out_item, "type") == "output_item.done" and _item_get(out_item, "phase") is not None: + latest_phase = _item_get(out_item, "phase") + break + + return resp, latest_phase +``` + +## Notes + +- Use `/v1/responses` for GPT Codex models. +- Preserve full assistant output history for best multi-turn behavior. +- If `phase` metadata is dropped during history reconstruction, output quality can degrade on long-running tasks. From 3e6c10a0710e7c64f27f5875bbf1a30702b4f143 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 24 Feb 2026 19:40:09 +0530 Subject: [PATCH 015/132] security: fix critical/high CVEs in OS-level libs and NPM transitive --- Dockerfile | 20 ++++++++-- ci_cd/security_scans.sh | 1 + docker/Dockerfile.custom_ui | 23 +++++++++++- docker/Dockerfile.database | 20 ++++++++-- docker/Dockerfile.dev | 37 ++++++++++++++++--- docker/Dockerfile.non_root | 23 +++++++++--- docs/my-website/package.json | 13 ++++++- litellm-js/spend-logs/package.json | 18 +++++++-- package.json | 18 +++++++-- requirements.txt | 4 +- tests/proxy_admin_ui_tests/package.json | 18 +++++++-- .../ui_unit_tests/package.json | 16 +++++++- ui/litellm-dashboard/package.json | 17 +++++++-- 13 files changed, 191 insertions(+), 37 deletions(-) diff --git a/Dockerfile b/Dockerfile index 5e93a0c627e..83d3640e763 100644 --- a/Dockerfile +++ b/Dockerfile @@ -49,7 +49,7 @@ USER root # Install runtime dependencies (libsndfile needed for audio processing on ARM64) RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip libsndfile && \ - npm install -g npm@latest tar@7.5.7 glob@11.1.0 @isaacs/brace-expansion@5.0.1 && \ + npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 && \ # SECURITY FIX: npm bundles tar, glob, and brace-expansion at multiple nested # levels inside its dependency tree. `npm install -g ` only creates a # SEPARATE global package, it does NOT replace npm's internal copies. @@ -64,6 +64,12 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip libsndfile find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ done && \ + find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done && \ + find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ + done && \ npm cache clean --force WORKDIR /app @@ -90,14 +96,20 @@ RUN find /usr/lib -type f -path "*/tornado/test/*" -delete && \ # npm with old vulnerable deps at /usr/lib/python3.*/site-packages/nodejs_wheel/. # Patch every copy of tar, glob, and brace-expansion inside that tree. RUN GLOBAL="$(npm root -g)" && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/tar" -type d | while read d; do \ + find /usr/lib -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ done && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/glob" -type d | while read d; do \ + find /usr/lib -type d -name "glob" -path "*/node_modules/glob" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \ done && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/@isaacs/brace-expansion" -type d | while read d; do \ + find /usr/lib -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ + done && \ + find /usr/lib -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done && \ + find /usr/lib -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done # Install semantic_router and aurelio-sdk using script diff --git a/ci_cd/security_scans.sh b/ci_cd/security_scans.sh index b60e7ec5235..0e50f15d043 100755 --- a/ci_cd/security_scans.sh +++ b/ci_cd/security_scans.sh @@ -160,6 +160,7 @@ run_grype_scans() { "CVE-2026-0775" # npm cli incorrect permission assignment - no fix available yet, npm is only used at build/prisma-generate time "GHSA-3ppc-4f35-3m26" # minimatch ReDoS via repeated wildcards - from nodejs_wheel bundled npm, not used in application runtime code "GHSA-83g3-92jg-28cx" # tar arbitrary file read/write via hardlink - from nodejs_wheel bundled npm, not used in application runtime code + "CVE-2026-25639" # axios - full fix requires 1.x major version bump; pinned to >=0.30.2 to clear other axios CVEs, upgrade to 1.x in follow-up ) # Build JSON array of allowlisted CVE IDs for jq diff --git a/docker/Dockerfile.custom_ui b/docker/Dockerfile.custom_ui index 177d7b7b12a..fb98846a6cc 100644 --- a/docker/Dockerfile.custom_ui +++ b/docker/Dockerfile.custom_ui @@ -5,8 +5,21 @@ FROM ghcr.io/berriai/litellm:litellm_fwd_server_root_path-dev WORKDIR /app # Install Node.js and npm (adjust version as needed) -RUN apt-get update && apt-get install -y nodejs npm && \ - npm install -g npm@latest tar@7.5.7 glob@11.1.0 @isaacs/brace-expansion@5.0.1 && \ +RUN apt-get update && apt-get upgrade -y \ + libxml2 \ + libexpat1 \ + openssl \ + libssl3 \ + git \ + libkrb5-3 \ + libglib2.0-0 \ + wget \ + libaom3 \ + libxslt1.1 \ + libgnutls30 \ + libc6 && \ + apt-get install -y nodejs npm && \ + npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 && \ GLOBAL="$(npm root -g)" && \ find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ @@ -17,6 +30,12 @@ RUN apt-get update && apt-get install -y nodejs npm && \ find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ done && \ + find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done && \ + find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ + done && \ npm cache clean --force # Copy the UI source into the container diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index a6fcd98ab6d..371766bd9db 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -50,7 +50,7 @@ USER root # Install runtime dependencies RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip libsndfile && \ - npm install -g npm@latest tar@7.5.7 glob@11.1.0 @isaacs/brace-expansion@5.0.1 && \ + npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 && \ GLOBAL="$(npm root -g)" && \ find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ @@ -61,6 +61,12 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip libsndfile find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ done && \ + find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done && \ + find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ + done && \ npm cache clean --force WORKDIR /app @@ -79,14 +85,20 @@ RUN pip install *.whl /wheels/* --no-index --find-links=/wheels/ && rm -f *.whl # npm with old vulnerable deps at /usr/lib/python3.*/site-packages/nodejs_wheel/. # Patch every copy of tar, glob, and brace-expansion inside that tree. RUN GLOBAL="$(npm root -g)" && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/tar" -type d | while read d; do \ + find /usr/lib -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ done && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/glob" -type d | while read d; do \ + find /usr/lib -type d -name "glob" -path "*/node_modules/glob" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \ done && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/@isaacs/brace-expansion" -type d | while read d; do \ + find /usr/lib -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ + done && \ + find /usr/lib -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done && \ + find /usr/lib -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done # Install semantic_router and aurelio-sdk using script diff --git a/docker/Dockerfile.dev b/docker/Dockerfile.dev index bc1d22d5e05..a5312dec9e3 100644 --- a/docker/Dockerfile.dev +++ b/docker/Dockerfile.dev @@ -56,13 +56,26 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root # Install only runtime dependencies -RUN apt-get update && apt-get install -y --no-install-recommends \ - libssl3 \ +RUN apt-get update && apt-get upgrade -y \ + libxml2 \ + libexpat1 \ + openssl \ + libssl3 \ + git \ + libkrb5-3 \ + libglib2.0-0 \ + wget \ + libaom3 \ + libxslt1.1 \ + libgnutls30 \ + libc6 \ + && apt-get install -y --no-install-recommends \ + libssl3 \ libatomic1 \ nodejs \ npm \ && rm -rf /var/lib/apt/lists/* \ - && npm install -g npm@latest tar@7.5.7 glob@11.1.0 @isaacs/brace-expansion@5.0.1 \ + && npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 \ && GLOBAL="$(npm root -g)" \ && find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ @@ -73,6 +86,12 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ && find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ done \ + && find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done \ + && find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ + done \ && npm cache clean --force WORKDIR /app @@ -95,14 +114,20 @@ RUN pip install --no-cache-dir *.whl /wheels/* --no-index --find-links=/wheels/ # npm with old vulnerable deps at /usr/lib/python3.*/site-packages/nodejs_wheel/. # Patch every copy of tar, glob, and brace-expansion inside that tree. RUN GLOBAL="$(npm root -g)" && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/tar" -type d | while read d; do \ + find /usr/lib -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ done && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/glob" -type d | while read d; do \ + find /usr/lib -type d -name "glob" -path "*/node_modules/glob" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \ done && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/@isaacs/brace-expansion" -type d | while read d; do \ + find /usr/lib -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ + done && \ + find /usr/lib -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done && \ + find /usr/lib -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done # Generate prisma client and set permissions diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 004377e19b3..fda591df083 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -80,7 +80,7 @@ ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \ XDG_CACHE_HOME=/app/.cache \ PATH="/usr/lib/python3.13/site-packages/nodejs/bin:${PATH}" -RUN pip install --no-cache-dir prisma==0.11.0 nodejs-wheel-binaries==24.12.0 \ +RUN pip install --no-cache-dir prisma==0.11.0 nodejs-wheel-binaries==24.13.1 \ && mkdir -p /app/.cache/npm RUN NPM_CONFIG_CACHE=/app/.cache/npm \ @@ -105,7 +105,8 @@ RUN for i in 1 2 3; do \ && for i in 1 2 3; do \ apk add --no-cache python3 py3-pip bash openssl tzdata nodejs npm supervisor && break || sleep 5; \ done \ - && npm install -g npm@latest tar@7.5.7 glob@11.1.0 @isaacs/brace-expansion@5.0.1 \ + && apk upgrade --no-cache nodejs \ + && npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 \ && GLOBAL="$(npm root -g)" \ && find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ @@ -116,6 +117,12 @@ RUN for i in 1 2 3; do \ && find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ done \ + && find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done \ + && find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ + done \ && npm cache clean --force # Copy artifacts from builder @@ -162,14 +169,20 @@ RUN pip install --no-index --find-links=/wheels/ -r requirements.txt && \ # npm with old vulnerable deps at /usr/lib/python3.*/site-packages/nodejs_wheel/. # Patch every copy of tar, glob, and brace-expansion inside that tree. RUN GLOBAL="$(npm root -g)" && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/tar" -type d | while read d; do \ + find /usr/lib -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ done && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/glob" -type d | while read d; do \ + find /usr/lib -type d -name "glob" -path "*/node_modules/glob" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \ done && \ - find /usr/lib -path "*/nodejs_wheel/*/node_modules/@isaacs/brace-expansion" -type d | while read d; do \ + find /usr/lib -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \ + done && \ + find /usr/lib -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \ + done && \ + find /usr/lib -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ + rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done # Permissions, cleanup, and Prisma prep diff --git a/docs/my-website/package.json b/docs/my-website/package.json index c4fa04c96a7..2dad6d0f16b 100644 --- a/docs/my-website/package.json +++ b/docs/my-website/package.json @@ -62,6 +62,8 @@ "gray-matter": "4.0.3", "glob": ">=11.1.0", "tar": ">=7.5.8", + "minimatch": ">=10.2.1", + "diff": ">=8.0.3", "@isaacs/brace-expansion": ">=5.0.1", "node-forge": ">=1.3.2", "mdast-util-to-hast": ">=13.2.1", @@ -81,6 +83,15 @@ "url-loader": { "ajv": "6.14.0" }, - "minimatch": "10.2.1" + "@babel/traverse": ">=7.23.2", + "ws": ">=7.5.10", + "http-proxy-middleware": ">=2.0.9", + "tar-fs": ">=2.1.4", + "webpack-dev-middleware": ">=5.3.4", + "braces": ">=3.0.3", + "axios": ">=0.30.2", + "webpack": ">=5.94.0", + "serve-static": ">=1.16.0", + "path-to-regexp": ">=0.1.12" } } \ No newline at end of file diff --git a/litellm-js/spend-logs/package.json b/litellm-js/spend-logs/package.json index 67292567145..5a7a08cb9ef 100644 --- a/litellm-js/spend-logs/package.json +++ b/litellm-js/spend-logs/package.json @@ -12,7 +12,19 @@ }, "overrides": { "glob": ">=11.1.0", - "tar": ">=7.5.7", - "@isaacs/brace-expansion": ">=5.0.1" + "tar": ">=7.5.8", + "minimatch": ">=10.2.1", + "diff": ">=8.0.3", + "@isaacs/brace-expansion": ">=5.0.1", + "@babel/traverse": ">=7.23.2", + "ws": ">=7.5.10", + "http-proxy-middleware": ">=2.0.9", + "tar-fs": ">=2.1.4", + "webpack-dev-middleware": ">=5.3.4", + "braces": ">=3.0.3", + "axios": ">=0.30.2", + "webpack": ">=5.94.0", + "serve-static": ">=1.16.0", + "path-to-regexp": ">=0.1.12" } -} +} \ No newline at end of file diff --git a/package.json b/package.json index ab9e15f46a7..a45e116b277 100644 --- a/package.json +++ b/package.json @@ -12,7 +12,19 @@ }, "overrides": { "glob": ">=11.1.0", - "tar": ">=7.5.7", - "@isaacs/brace-expansion": ">=5.0.1" + "tar": ">=7.5.8", + "minimatch": ">=10.2.1", + "diff": ">=8.0.3", + "@isaacs/brace-expansion": ">=5.0.1", + "@babel/traverse": ">=7.23.2", + "ws": ">=7.5.10", + "http-proxy-middleware": ">=2.0.9", + "tar-fs": ">=2.1.4", + "webpack-dev-middleware": ">=5.3.4", + "braces": ">=3.0.3", + "axios": ">=0.30.2", + "webpack": ">=5.94.0", + "serve-static": ">=1.16.0", + "path-to-regexp": ">=0.1.12" } -} +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 832c3f75b72..dbde6ababc9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,8 @@ urllib3>=2.6.0 # CVE-2025-66471, CVE-2025-66418, CVE-2026-21441 tornado>=6.5.3 # CVE-2025-67725, CVE-2025-67726, CVE-2025-67724 filelock>=3.20.1 # CVE-2025-68146 +h11>=0.16.0 # CVE-2025-43859, GHSA-vqfr-h8mv-ghfj — HTTP request smuggling +wheel>=0.46.2 # CVE-2026-24049 — path traversal Pillow==12.1.1 #GHSA-cfh3-3jmp-rvhc cryptography==46.0.5 #GHSA-r6ph-v2qm-q3c2 @@ -21,7 +23,7 @@ boto3==1.40.53 # aws bedrock/sagemaker calls (has bedrock-agentcore-control, com redis==5.2.1 # redis caching redisvl==0.4.1 ## redis semantic caching prisma==0.11.0 # for db -nodejs-wheel-binaries==24.12.0 ## required by prisma for migrations, prevents runtime download (updated from nodejs-bin for security fixes) +nodejs-wheel-binaries==24.13.1 ## required by prisma for migrations, prevents runtime download (updated from nodejs-bin for security fixes) mangum==0.17.0 # for aws lambda functions pynacl==1.6.2 # for encrypting keys google-cloud-aiplatform==1.133.0 # for vertex ai calls diff --git a/tests/proxy_admin_ui_tests/package.json b/tests/proxy_admin_ui_tests/package.json index 48de2c1dba2..037ec7082a7 100644 --- a/tests/proxy_admin_ui_tests/package.json +++ b/tests/proxy_admin_ui_tests/package.json @@ -13,7 +13,19 @@ }, "overrides": { "glob": ">=11.1.0", - "tar": ">=7.5.7", - "@isaacs/brace-expansion": ">=5.0.1" + "tar": ">=7.5.8", + "minimatch": ">=10.2.1", + "diff": ">=8.0.3", + "@isaacs/brace-expansion": ">=5.0.1", + "@babel/traverse": ">=7.23.2", + "ws": ">=7.5.10", + "http-proxy-middleware": ">=2.0.9", + "tar-fs": ">=2.1.4", + "webpack-dev-middleware": ">=5.3.4", + "braces": ">=3.0.3", + "axios": ">=0.30.2", + "webpack": ">=5.94.0", + "serve-static": ">=1.16.0", + "path-to-regexp": ">=0.1.12" } -} +} \ No newline at end of file diff --git a/tests/proxy_admin_ui_tests/ui_unit_tests/package.json b/tests/proxy_admin_ui_tests/ui_unit_tests/package.json index 4c7d7addf0e..eb9c7473a5b 100644 --- a/tests/proxy_admin_ui_tests/ui_unit_tests/package.json +++ b/tests/proxy_admin_ui_tests/ui_unit_tests/package.json @@ -25,7 +25,19 @@ }, "overrides": { "glob": ">=11.1.0", - "tar": ">=7.5.7", - "@isaacs/brace-expansion": ">=5.0.1" + "tar": ">=7.5.8", + "minimatch": ">=10.2.1", + "diff": ">=8.0.3", + "@isaacs/brace-expansion": ">=5.0.1", + "@babel/traverse": ">=7.23.2", + "ws": ">=7.5.10", + "http-proxy-middleware": ">=2.0.9", + "tar-fs": ">=2.1.4", + "webpack-dev-middleware": ">=5.3.4", + "braces": ">=3.0.3", + "axios": ">=0.30.2", + "webpack": ">=5.94.0", + "serve-static": ">=1.16.0", + "path-to-regexp": ">=0.1.12" } } \ No newline at end of file diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index b05d707d5ab..567673c0989 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -86,14 +86,25 @@ "mermaid": ">=11.10.0", "js-yaml": ">=4.1.1", "glob": ">=11.1.0", - "tar": ">=7.5.7", + "tar": ">=7.5.8", + "minimatch": ">=10.2.1", "@isaacs/brace-expansion": ">=5.0.1", "node-forge": ">=1.3.2", "lodash-es": ">=4.17.23", - "lodash": ">=4.17.23" + "lodash": ">=4.17.23", + "@babel/traverse": ">=7.23.2", + "ws": ">=7.5.10", + "http-proxy-middleware": ">=2.0.9", + "tar-fs": ">=2.1.4", + "webpack-dev-middleware": ">=5.3.4", + "braces": ">=3.0.3", + "axios": ">=0.30.2", + "webpack": ">=5.94.0", + "serve-static": ">=1.16.0", + "path-to-regexp": ">=0.1.12" }, "engines": { "node": ">=18.17.0", "npm": ">=8.3.0" } -} +} \ No newline at end of file From 4652c73259e2e3b1d3cadea360f2bc5f4b27d1eb Mon Sep 17 00:00:00 2001 From: Sean Marsh Glover Date: Tue, 24 Feb 2026 11:16:59 -0500 Subject: [PATCH 016/132] feat(proxy): limit concurrent health checks with health_check_concurrency (#20584) * staged first pass * black * Update litellm/proxy/health_check.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * simpler * restore cached logo * fix tests for perform_health_check max_concurrency arg * implement pr suggestion * and the helm chart * add configureable resources and probes to the deployment in the helm chart * more helm chart unittests * move some background healthcheck loggin to debug --------- Co-authored-by: Sean Glover Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- deploy/charts/litellm-helm/README.md | 4 + .../templates/configmap-litellm.yaml | 2 +- .../litellm-helm/templates/deployment.yaml | 25 +- .../litellm-helm/tests/deployment_tests.yaml | 148 ++++++++++- deploy/charts/litellm-helm/values.yaml | 25 ++ litellm/proxy/_types.py | 7 + litellm/proxy/health_check.py | 194 ++++++++++++-- .../shared_health_check_manager.py | 96 +++---- .../health_endpoints/_health_endpoints.py | 114 ++++---- litellm/proxy/proxy_server.py | 247 ++++++++++++++---- .../litellm_utils_tests/test_health_check.py | 119 ++++++++- tests/proxy_unit_tests/test_proxy_server.py | 8 +- 12 files changed, 809 insertions(+), 180 deletions(-) diff --git a/deploy/charts/litellm-helm/README.md b/deploy/charts/litellm-helm/README.md index 2fa856843f3..74e70f4aeb4 100644 --- a/deploy/charts/litellm-helm/README.md +++ b/deploy/charts/litellm-helm/README.md @@ -36,6 +36,10 @@ If `db.useStackgresOperator` is used (not yet implemented): | `serviceAccount.create` | Whether or not to create a Kubernetes Service Account for this deployment. The default is `false` because LiteLLM has no need to access the Kubernetes API. | `false` | | `service.type` | Kubernetes Service type (e.g. `LoadBalancer`, `ClusterIP`, etc.) | `ClusterIP` | | `service.port` | TCP port that the Kubernetes Service will listen on. Also the TCP port within the Pod that the proxy will listen on. | `4000` | +| `livenessProbe.*` | Liveness probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` | +| `readinessProbe.*` | Readiness probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` | +| `startupProbe.*` | Startup probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` | +| `resources.*` | CPU/memory requests and limits for the LiteLLM container. | `{}` | | `service.loadBalancerClass` | Optional LoadBalancer implementation class (only used when `service.type` is `LoadBalancer`) | `""` | | `ingress.labels` | Additional labels for the Ingress resource | `{}` | | `ingress.*` | See [values.yaml](./values.yaml) for example settings | N/A | diff --git a/deploy/charts/litellm-helm/templates/configmap-litellm.yaml b/deploy/charts/litellm-helm/templates/configmap-litellm.yaml index cf35917da03..acbe4e3a4b5 100644 --- a/deploy/charts/litellm-helm/templates/configmap-litellm.yaml +++ b/deploy/charts/litellm-helm/templates/configmap-litellm.yaml @@ -6,4 +6,4 @@ metadata: data: config.yaml: | {{ .Values.proxy_config | toYaml | indent 6 }} -{{- end }} \ No newline at end of file +{{- end }} diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index 4ac5582d060..df483ab927d 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -158,18 +158,31 @@ spec: {{- end }} livenessProbe: httpGet: - path: /health/liveliness + path: {{ .Values.livenessProbe.path | quote }} port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }} + initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds }} + periodSeconds: {{ .Values.livenessProbe.periodSeconds }} + timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds }} + successThreshold: {{ .Values.livenessProbe.successThreshold }} + failureThreshold: {{ .Values.livenessProbe.failureThreshold }} readinessProbe: httpGet: - path: /health/readiness + path: {{ .Values.readinessProbe.path | quote }} port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }} + initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds }} + periodSeconds: {{ .Values.readinessProbe.periodSeconds }} + timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds }} + successThreshold: {{ .Values.readinessProbe.successThreshold }} + failureThreshold: {{ .Values.readinessProbe.failureThreshold }} startupProbe: httpGet: - path: /health/readiness + path: {{ .Values.startupProbe.path | quote }} port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }} - failureThreshold: 30 - periodSeconds: 10 + initialDelaySeconds: {{ .Values.startupProbe.initialDelaySeconds }} + periodSeconds: {{ .Values.startupProbe.periodSeconds }} + timeoutSeconds: {{ .Values.startupProbe.timeoutSeconds }} + successThreshold: {{ .Values.startupProbe.successThreshold }} + failureThreshold: {{ .Values.startupProbe.failureThreshold }} resources: {{- toYaml .Values.resources | nindent 12 }} volumeMounts: @@ -235,4 +248,4 @@ spec: {{- if .Values.topologySpreadConstraints }} topologySpreadConstraints: {{- toYaml .Values.topologySpreadConstraints | nindent 8 }} - {{- end }} \ No newline at end of file + {{- end }} diff --git a/deploy/charts/litellm-helm/tests/deployment_tests.yaml b/deploy/charts/litellm-helm/tests/deployment_tests.yaml index f1229e10235..2e9c48043de 100644 --- a/deploy/charts/litellm-helm/tests/deployment_tests.yaml +++ b/deploy/charts/litellm-helm/tests/deployment_tests.yaml @@ -159,4 +159,150 @@ tests: value: -c - equal: path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[2] - value: echo "Container stopping" \ No newline at end of file + value: echo "Container stopping" + - it: should render background health check settings from proxy_config.general_settings + template: configmap-litellm.yaml + set: + proxy_config.general_settings.background_health_checks: true + proxy_config.general_settings.health_check_interval: 240 + proxy_config.general_settings.health_check_concurrency: 16 + proxy_config.general_settings.health_check_details: false + asserts: + - matchRegex: + path: data["config.yaml"] + pattern: '(?m)^\s*background_health_checks:\s*true$' + - matchRegex: + path: data["config.yaml"] + pattern: '(?m)^\s*health_check_interval:\s*240$' + - matchRegex: + path: data["config.yaml"] + pattern: '(?m)^\s*health_check_concurrency:\s*16$' + - matchRegex: + path: data["config.yaml"] + pattern: '(?m)^\s*health_check_details:\s*false$' + - it: should allow overriding liveness, readiness, and startup probes + template: deployment.yaml + set: + livenessProbe: + path: /custom/livez + initialDelaySeconds: 5 + periodSeconds: 15 + timeoutSeconds: 5 + successThreshold: 1 + failureThreshold: 5 + readinessProbe: + path: /custom/readyz + initialDelaySeconds: 10 + periodSeconds: 20 + timeoutSeconds: 6 + successThreshold: 1 + failureThreshold: 6 + startupProbe: + path: /custom/startupz + initialDelaySeconds: 15 + periodSeconds: 25 + timeoutSeconds: 7 + successThreshold: 1 + failureThreshold: 40 + asserts: + - equal: + path: spec.template.spec.containers[0].livenessProbe.httpGet.path + value: /custom/livez + - equal: + path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds + value: 5 + - equal: + path: spec.template.spec.containers[0].readinessProbe.httpGet.path + value: /custom/readyz + - equal: + path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds + value: 6 + - equal: + path: spec.template.spec.containers[0].startupProbe.httpGet.path + value: /custom/startupz + - equal: + path: spec.template.spec.containers[0].startupProbe.failureThreshold + value: 40 + - it: should render container resources from values + template: deployment.yaml + set: + resources: + limits: + cpu: 500m + memory: 2Gi + requests: + cpu: 250m + memory: 1Gi + asserts: + - equal: + path: spec.template.spec.containers[0].resources.limits.cpu + value: 500m + - equal: + path: spec.template.spec.containers[0].resources.limits.memory + value: 2Gi + - equal: + path: spec.template.spec.containers[0].resources.requests.cpu + value: 250m + - equal: + path: spec.template.spec.containers[0].resources.requests.memory + value: 1Gi + - it: should keep default probes and empty resources unchanged + template: deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].livenessProbe.httpGet.path + value: /health/liveliness + - equal: + path: spec.template.spec.containers[0].livenessProbe.initialDelaySeconds + value: 0 + - equal: + path: spec.template.spec.containers[0].livenessProbe.periodSeconds + value: 10 + - equal: + path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds + value: 1 + - equal: + path: spec.template.spec.containers[0].livenessProbe.successThreshold + value: 1 + - equal: + path: spec.template.spec.containers[0].livenessProbe.failureThreshold + value: 3 + - equal: + path: spec.template.spec.containers[0].readinessProbe.httpGet.path + value: /health/readiness + - equal: + path: spec.template.spec.containers[0].readinessProbe.initialDelaySeconds + value: 0 + - equal: + path: spec.template.spec.containers[0].readinessProbe.periodSeconds + value: 10 + - equal: + path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds + value: 1 + - equal: + path: spec.template.spec.containers[0].readinessProbe.successThreshold + value: 1 + - equal: + path: spec.template.spec.containers[0].readinessProbe.failureThreshold + value: 3 + - equal: + path: spec.template.spec.containers[0].startupProbe.httpGet.path + value: /health/readiness + - equal: + path: spec.template.spec.containers[0].startupProbe.initialDelaySeconds + value: 0 + - equal: + path: spec.template.spec.containers[0].startupProbe.periodSeconds + value: 10 + - equal: + path: spec.template.spec.containers[0].startupProbe.timeoutSeconds + value: 1 + - equal: + path: spec.template.spec.containers[0].startupProbe.successThreshold + value: 1 + - equal: + path: spec.template.spec.containers[0].startupProbe.failureThreshold + value: 30 + - equal: + path: spec.template.spec.containers[0].resources + value: {} diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index cea25974bb0..d62f5b29c2b 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -84,6 +84,31 @@ service: separateHealthApp: false separateHealthPort: 8081 +# Probe tuning for proxy container +livenessProbe: + path: /health/liveliness + initialDelaySeconds: 0 + periodSeconds: 10 + timeoutSeconds: 1 + successThreshold: 1 + failureThreshold: 3 + +readinessProbe: + path: /health/readiness + initialDelaySeconds: 0 + periodSeconds: 10 + timeoutSeconds: 1 + successThreshold: 1 + failureThreshold: 3 + +startupProbe: + path: /health/readiness + initialDelaySeconds: 0 + periodSeconds: 10 + timeoutSeconds: 1 + successThreshold: 1 + failureThreshold: 30 + ingress: enabled: false className: "nginx" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 95739834a9a..f354e28acd7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2079,6 +2079,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): health_check_interval: int = Field( 300, description="background health check interval in seconds" ) + health_check_concurrency: Optional[int] = Field( + None, + description=( + "limit concurrent health checks per cycle; when unset, " + "health checks run without a concurrency cap" + ), + ) alerting: Optional[List] = Field( None, description="List of alerting integrations. Today, just slack - `alerting: ['slack']`", diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 427a16a9801..4777f64405d 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -3,6 +3,9 @@ import asyncio import logging import random +import sys +import threading +import time from typing import List, Optional import litellm @@ -23,6 +26,29 @@ ILLEGAL_DISPLAY_PARAMS = [ MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"] +def _get_process_rss_mb() -> Optional[float]: + """ + Get process RSS memory in MB. + On Linux, ru_maxrss is in KB. On macOS, ru_maxrss is in bytes. + """ + try: + import resource + + ru_maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if sys.platform == "darwin": + return float(ru_maxrss) / (1024 * 1024) + return float(ru_maxrss) / 1024 + except Exception: + return None + + +def _rss_mb_for_log() -> str: + rss_mb = _get_process_rss_mb() + if rss_mb is None: + return "unknown" + return f"{rss_mb:.2f}" + + def _get_random_llm_message(): """ Get a random message from the LLM. @@ -67,26 +93,29 @@ async def run_with_timeout(task, timeout): try: return await asyncio.wait_for(task, timeout) except asyncio.TimeoutError: - task.cancel() - # Only cancel child tasks of the current task - current_task = asyncio.current_task() - for t in asyncio.all_tasks(): - if t != current_task: - t.cancel() - try: - await asyncio.wait_for(task, 0.1) # Give 100ms for cleanup - except (asyncio.TimeoutError, asyncio.CancelledError, Exception): - pass + # `asyncio.wait_for()` already cancels only the awaited task on timeout. + # Do not cancel unrelated sibling health check tasks. return {"error": "Timeout exceeded"} -async def _perform_health_check(model_list: list, details: Optional[bool] = True): +async def _perform_health_check( + model_list: list, + details: Optional[bool] = True, + max_concurrency: Optional[int] = None, + instrumentation_context: Optional[dict] = None, +): """ Perform a health check for each model in the list. + + max_concurrency: Optional limit on concurrent health check requests. """ - tasks = [] - for model in model_list: + instrumentation_context = instrumentation_context or {} + instrumentation_enabled = bool(instrumentation_context.get("enabled", False)) + cycle_id = instrumentation_context.get("cycle_id", "unknown") + source = instrumentation_context.get("source", "unknown") + + async def _run_model_health_check(model: dict): litellm_params = model["litellm_params"] model_info = model.get("model_info", {}) mode = model_info.get("mode", None) @@ -95,9 +124,9 @@ async def _perform_health_check(model_list: list, details: Optional[bool] = True ) timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS - task = run_with_timeout( + return await run_with_timeout( litellm.ahealth_check( - model["litellm_params"], + litellm_params, mode=mode, prompt=DEFAULT_HEALTH_CHECK_PROMPT, input=["test from litellm"], @@ -105,9 +134,73 @@ async def _perform_health_check(model_list: list, details: Optional[bool] = True timeout, ) - tasks.append(task) + async def _run_health_checks_with_bounded_concurrency( + models: list, concurrency_limit: int + ) -> tuple[list, int]: + """ + Run health checks with at most `concurrency_limit` active tasks. + Preserves result ordering to match `models`. + """ + results: list = [None] * len(models) + tasks_to_index: dict[asyncio.Task, int] = {} + model_iter = iter(enumerate(models)) + peak_in_flight = 0 - results = await asyncio.gather(*tasks, return_exceptions=True) + def _schedule_next() -> bool: + nonlocal peak_in_flight + try: + idx, next_model = next(model_iter) + except StopIteration: + return False + task = asyncio.create_task(_run_model_health_check(next_model)) + tasks_to_index[task] = idx + peak_in_flight = max(peak_in_flight, len(tasks_to_index)) + return True + + for _ in range(min(concurrency_limit, len(models))): + _schedule_next() + + while tasks_to_index: + done, _ = await asyncio.wait( + set(tasks_to_index.keys()), + return_when=asyncio.FIRST_COMPLETED, + ) + for task in done: + idx = tasks_to_index.pop(task) + try: + results[idx] = task.result() + except Exception as e: + results[idx] = e + _schedule_next() + + return results, peak_in_flight + + dispatch_mode = "unbounded" + peak_in_flight = 0 + if isinstance(max_concurrency, int) and max_concurrency > 0: + dispatch_mode = "bounded" + results, peak_in_flight = await _run_health_checks_with_bounded_concurrency( + model_list, max_concurrency + ) + else: + tasks = [ + asyncio.create_task(_run_model_health_check(model)) for model in model_list + ] + peak_in_flight = len(tasks) + results = await asyncio.gather(*tasks, return_exceptions=True) + + if instrumentation_enabled: + logger.debug( + "health_check_dispatch_summary source=%s cycle_id=%s mode=%s model_count=%d max_concurrency=%s peak_in_flight=%d thread_count=%d rss_mb=%s", + source, + cycle_id, + dispatch_mode, + len(model_list), + max_concurrency, + peak_in_flight, + threading.active_count(), + _rss_mb_for_log(), + ) healthy_endpoints = [] unhealthy_endpoints = [] @@ -190,6 +283,8 @@ async def perform_health_check( model: Optional[str] = None, cli_model: Optional[str] = None, details: Optional[bool] = True, + max_concurrency: Optional[int] = None, + instrumentation_context: Optional[dict] = None, ): """ Perform a health check on the system. @@ -197,14 +292,28 @@ async def perform_health_check( Returns: (bool): True if the health check passes, False otherwise. """ + instrumentation_context = instrumentation_context or {} + instrumentation_enabled = bool(instrumentation_context.get("enabled", False)) + cycle_id = instrumentation_context.get("cycle_id", "unknown") + source = instrumentation_context.get("source", "unknown") + if not model_list: if cli_model: model_list = [ {"model_name": cli_model, "litellm_params": {"model": cli_model}} ] else: + if instrumentation_enabled: + logger.debug( + "health_check_cycle_skipped source=%s cycle_id=%s reason=no_models", + source, + cycle_id, + ) return [], [] + cycle_start_time = time.monotonic() + requested_model_count = len(model_list) + if model is not None: _new_model_list = [ x for x in model_list if x["litellm_params"]["model"] == model @@ -213,11 +322,56 @@ async def perform_health_check( _new_model_list = [x for x in model_list if x["model_name"] == model] model_list = _new_model_list + post_filter_model_count = len(model_list) model_list = filter_deployments_by_id( model_list=model_list ) # filter duplicate deployments (e.g. when model alias'es are used) - healthy_endpoints, unhealthy_endpoints = await _perform_health_check( - model_list, details - ) + deduped_model_count = len(model_list) + + if instrumentation_enabled: + logger.debug( + "health_check_cycle_start source=%s cycle_id=%s requested_model_count=%d post_model_filter_count=%d deduped_model_count=%d max_concurrency=%s thread_count=%d rss_mb=%s", + source, + cycle_id, + requested_model_count, + post_filter_model_count, + deduped_model_count, + max_concurrency, + threading.active_count(), + _rss_mb_for_log(), + ) + + try: + healthy_endpoints, unhealthy_endpoints = await _perform_health_check( + model_list, + details, + max_concurrency=max_concurrency, + instrumentation_context=instrumentation_context, + ) + except Exception: + if instrumentation_enabled: + logger.exception( + "health_check_cycle_failed source=%s cycle_id=%s model_count=%d duration_ms=%.2f thread_count=%d rss_mb=%s", + source, + cycle_id, + deduped_model_count, + (time.monotonic() - cycle_start_time) * 1000, + threading.active_count(), + _rss_mb_for_log(), + ) + raise + + if instrumentation_enabled: + logger.debug( + "health_check_cycle_complete source=%s cycle_id=%s model_count=%d healthy_count=%d unhealthy_count=%d duration_ms=%.2f thread_count=%d rss_mb=%s", + source, + cycle_id, + deduped_model_count, + len(healthy_endpoints), + len(unhealthy_endpoints), + (time.monotonic() - cycle_start_time) * 1000, + threading.active_count(), + _rss_mb_for_log(), + ) return healthy_endpoints, unhealthy_endpoints diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index d0c99d84e94..ae18a42c02b 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -16,7 +16,7 @@ from litellm.proxy.health_check import perform_health_check class SharedHealthCheckManager: """ Manager for coordinating health checks across multiple pods using Redis. - + This class implements a shared health check state mechanism that: - Prevents duplicate health checks across pods - Caches health check results with configurable TTL @@ -58,7 +58,7 @@ class SharedHealthCheckManager: async def acquire_health_check_lock(self) -> bool: """ Attempt to acquire the global health check lock. - + Returns: bool: True if lock was acquired, False otherwise """ @@ -74,7 +74,7 @@ class SharedHealthCheckManager: nx=True, # Only set if key doesn't exist ttl=self.lock_ttl, ) - + if acquired: verbose_proxy_logger.info( "Pod %s acquired health check lock", self.pod_id @@ -83,12 +83,10 @@ class SharedHealthCheckManager: verbose_proxy_logger.debug( "Pod %s failed to acquire health check lock", self.pod_id ) - + return acquired except Exception as e: - verbose_proxy_logger.error( - "Error acquiring health check lock: %s", str(e) - ) + verbose_proxy_logger.error("Error acquiring health check lock: %s", str(e)) return False async def release_health_check_lock(self) -> None: @@ -106,14 +104,12 @@ class SharedHealthCheckManager: "Pod %s released health check lock", self.pod_id ) except Exception as e: - verbose_proxy_logger.error( - "Error releasing health check lock: %s", str(e) - ) + verbose_proxy_logger.error("Error releasing health check lock: %s", str(e)) async def get_cached_health_check_results(self) -> Optional[Dict[str, Any]]: """ Get cached health check results from Redis. - + Returns: Optional[Dict]: Cached health check results or None if not found/expired """ @@ -123,7 +119,7 @@ class SharedHealthCheckManager: try: cache_key = self.get_health_check_cache_key() cached_data = await self.redis_cache.async_get_cache(cache_key) - + if cached_data is None: return None @@ -136,7 +132,7 @@ class SharedHealthCheckManager: # Check if the cache is still valid cache_timestamp = cached_results.get("timestamp", 0) current_time = time.time() - + if current_time - cache_timestamp > self.health_check_ttl: verbose_proxy_logger.debug("Cached health check results expired") return None @@ -151,13 +147,13 @@ class SharedHealthCheckManager: return None async def cache_health_check_results( - self, - healthy_endpoints: List[Dict[str, Any]], - unhealthy_endpoints: List[Dict[str, Any]] + self, + healthy_endpoints: List[Dict[str, Any]], + unhealthy_endpoints: List[Dict[str, Any]], ) -> None: """ Cache health check results in Redis. - + Args: healthy_endpoints: List of healthy endpoints unhealthy_endpoints: List of unhealthy endpoints @@ -181,7 +177,7 @@ class SharedHealthCheckManager: safe_dumps(cache_data), ttl=self.health_check_ttl, ) - + verbose_proxy_logger.info( "Cached health check results for %d healthy and %d unhealthy endpoints", len(healthy_endpoints), @@ -189,29 +185,29 @@ class SharedHealthCheckManager: ) except Exception as e: - verbose_proxy_logger.error( - "Error caching health check results: %s", str(e) - ) + verbose_proxy_logger.error("Error caching health check results: %s", str(e)) async def perform_shared_health_check( - self, - model_list: List[Dict[str, Any]], - details: bool = True + self, + model_list: List[Dict[str, Any]], + details: bool = True, + max_concurrency: Optional[int] = None, ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: """ Perform health check with shared state coordination. - + This method: 1. First checks if there are recent cached results 2. If no recent cache, tries to acquire lock to run health check 3. If lock acquired, runs health check and caches results 4. If lock not acquired, waits briefly and tries to get cached results again 5. Falls back to running health check locally if no cache available - + Args: model_list: List of models to check details: Whether to include detailed information - + max_concurrency: Optional limit on concurrent health check requests + Returns: Tuple of (healthy_endpoints, unhealthy_endpoints) """ @@ -225,27 +221,29 @@ class SharedHealthCheckManager: # No recent cache, try to acquire lock lock_acquired = await self.acquire_health_check_lock() - + if lock_acquired: try: # We have the lock, run health check verbose_proxy_logger.info( - "Pod %s running health check for %d models", - self.pod_id, - len(model_list) + "Pod %s running health check for %d models", + self.pod_id, + len(model_list), ) - + healthy_endpoints, unhealthy_endpoints = await perform_health_check( - model_list=model_list, details=details + model_list=model_list, + details=details, + max_concurrency=max_concurrency, ) - + # Cache the results await self.cache_health_check_results( healthy_endpoints, unhealthy_endpoints ) - + return healthy_endpoints, unhealthy_endpoints - + finally: # Always release the lock await self.release_health_check_lock() @@ -254,10 +252,10 @@ class SharedHealthCheckManager: verbose_proxy_logger.debug( "Pod %s waiting for other pod to complete health check", self.pod_id ) - + # Wait a bit for the other pod to complete await asyncio.sleep(2) - + # Try to get cached results again cached_results = await self.get_cached_health_check_results() if cached_results is not None: @@ -265,19 +263,23 @@ class SharedHealthCheckManager: cached_results.get("healthy_endpoints", []), cached_results.get("unhealthy_endpoints", []), ) - + # Still no cache, fall back to local health check verbose_proxy_logger.warning( - "Pod %s falling back to local health check (no cache available)", - self.pod_id + "Pod %s falling back to local health check (no cache available)", + self.pod_id, + ) + + return await perform_health_check( + model_list=model_list, + details=details, + max_concurrency=max_concurrency, ) - - return await perform_health_check(model_list=model_list, details=details) async def is_health_check_in_progress(self) -> bool: """ Check if a health check is currently in progress by another pod. - + Returns: bool: True if health check is in progress, False otherwise """ @@ -297,7 +299,7 @@ class SharedHealthCheckManager: async def get_health_check_status(self) -> Dict[str, Any]: """ Get the current status of health check coordination. - + Returns: Dict containing status information """ @@ -320,7 +322,9 @@ class SharedHealthCheckManager: cached_results = await self.get_cached_health_check_results() status["cache_available"] = cached_results is not None if cached_results: - status["cache_age_seconds"] = time.time() - cached_results.get("timestamp", 0) + status["cache_age_seconds"] = time.time() - cached_results.get( + "timestamp", 0 + ) status["last_checked_by"] = cached_results.get("checked_by") except Exception as e: diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index da90696ec2d..3570b2dd6aa 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -110,26 +110,31 @@ def _resolve_os_environ_variables(params: dict) -> dict: def get_callback_identifier(callback): """ Get the callback identifier string, handling both strings and objects. - + This function extracts a string identifier from a callback, which can be: - A string (returned as-is) - An object with a callback_name attribute - An object registered in CustomLoggerRegistry - Falls back to callback_name() helper function - + Args: callback: The callback to identify (can be str or object) - + Returns: str: The callback identifier string """ if isinstance(callback, str): return callback - if hasattr(callback, 'callback_name') and callback.callback_name: + if hasattr(callback, "callback_name") and callback.callback_name: return callback.callback_name - if hasattr(callback, '__class__'): - callback_strs = CustomLoggerRegistry.get_all_callback_strs_from_class_type(callback.__class__) - if hasattr(callback, 'callback_name') and callback.callback_name in callback_strs: + if hasattr(callback, "__class__"): + callback_strs = CustomLoggerRegistry.get_all_callback_strs_from_class_type( + callback.__class__ + ) + if ( + hasattr(callback, "callback_name") + and callback.callback_name in callback_strs + ): return callback.callback_name if callback_strs: return callback_strs[0] @@ -151,7 +156,7 @@ services = Union[ "datadog_llm_observability", "generic_api", "arize", - "sqs" + "sqs", ], str, ] @@ -224,7 +229,7 @@ async def health_services_endpoint( # noqa: PLR0915 "datadog_llm_observability", "generic_api", "arize", - "sqs" + "sqs", ]: raise HTTPException( status_code=400, @@ -238,14 +243,14 @@ async def health_services_endpoint( # noqa: PLR0915 service_in_success_callbacks = True else: for cb in litellm.success_callback: - if hasattr(cb, 'callback_name') and cb.callback_name == service: + if hasattr(cb, "callback_name") and cb.callback_name == service: service_in_success_callbacks = True break cb_id = get_callback_identifier(cb) if cb_id == service: service_in_success_callbacks = True break - + if ( service == "openmeter" or service == "braintrust" @@ -320,6 +325,7 @@ async def health_services_endpoint( # noqa: PLR0915 ) elif service == "sqs": from litellm.integrations.sqs import SQSLogger + sqs_logger = SQSLogger() response = await sqs_logger.async_health_check() return { @@ -518,12 +524,12 @@ async def _save_health_check_to_db( def _build_model_param_to_info_mapping(model_list: list) -> dict: """ Build a mapping from model parameter to model info (model_name, model_id). - + Multiple models might share the same model parameter, so we use a list. - + Args: model_list: List of model configurations - + Returns: Dictionary mapping model parameter to list of model info dicts """ @@ -534,14 +540,16 @@ def _build_model_param_to_info_mapping(model_list: list) -> dict: model_id = model_info.get("id") litellm_params = model.get("litellm_params", {}) model_param = litellm_params.get("model") - + if model_param and model_name: if model_param not in model_param_to_info: model_param_to_info[model_param] = [] - model_param_to_info[model_param].append({ - "model_name": model_name, - "model_id": model_id, - }) + model_param_to_info[model_param].append( + { + "model_name": model_name, + "model_id": model_id, + } + ) return model_param_to_info @@ -552,19 +560,19 @@ def _aggregate_health_check_results( ) -> dict: """ Aggregate health check results per unique model. - + Uses (model_id, model_name) as key, or (None, model_name) if model_id is None. - + Args: model_param_to_info: Mapping from model parameter to model info healthy_endpoints: List of healthy endpoint results unhealthy_endpoints: List of unhealthy endpoint results - + Returns: Dictionary mapping (model_id, model_name) to aggregated health check results """ model_results = {} - + # Process healthy endpoints for endpoint in healthy_endpoints: model_param = endpoint.get("model") @@ -580,7 +588,7 @@ def _aggregate_health_check_results( "error_message": None, } model_results[key]["healthy_count"] += 1 - + # Process unhealthy endpoints for endpoint in unhealthy_endpoints: model_param = endpoint.get("model") @@ -600,7 +608,7 @@ def _aggregate_health_check_results( # Use the first error message encountered if not model_results[key]["error_message"] and error_message: model_results[key]["error_message"] = str(error_message)[:500] - + return model_results @@ -613,14 +621,14 @@ async def _save_health_check_results_if_changed( ): """ Save health check results to database, but only if status changed or >1 hour since last save. - + OPTIMIZATION: Only saves to database if the status has changed from the last saved check. This dramatically reduces database writes when health status remains stable. - + - Stable systems: ~1 write/hour per model (instead of 12 writes/hour with 5-min intervals) - Status changes: Immediate write (no delay) - Result: ~92% reduction in DB writes for stable systems, while maintaining real-time updates on changes - + Args: prisma_client: Database client model_results: Dictionary of aggregated health check results per model @@ -630,7 +638,7 @@ async def _save_health_check_results_if_changed( """ for result in model_results.values(): new_status = "healthy" if result["healthy_count"] > 0 else "unhealthy" - + # Check if we should save this result should_save = True lookup_key = result["model_id"] if result["model_id"] else result["model_name"] @@ -641,6 +649,7 @@ async def _save_health_check_results_if_changed( # Check if last check was recent (within 1 hour) if last_check.checked_at: from datetime import datetime, timezone + time_since_last_check = ( datetime.now(timezone.utc) - last_check.checked_at ).total_seconds() @@ -648,7 +657,7 @@ async def _save_health_check_results_if_changed( # This ensures we still get periodic updates even if status is stable if time_since_last_check < 3600: # 1 hour threshold should_save = False - + if should_save: asyncio.create_task( prisma_client.save_health_check_result( @@ -675,27 +684,27 @@ async def _save_background_health_checks_to_db( ): """ Save background health check results to database for each model. - + Maps health check endpoints back to their original models to get model_name and model_id. Aggregates results per unique model (by model_id if available, otherwise model_name). - + OPTIMIZATION: Only saves to database if the status has changed from the last saved check. This dramatically reduces database writes when health status remains stable. """ if prisma_client is None: return - + try: # Step 1: Build mapping from model parameter to model info model_param_to_info = _build_model_param_to_info_mapping(model_list) - + # Step 2: Aggregate health check results per unique model model_results = _aggregate_health_check_results( model_param_to_info, healthy_endpoints, unhealthy_endpoints, ) - + # Step 3: Get latest health checks for all models in one query to compare status latest_checks = await prisma_client.get_all_latest_health_checks() latest_checks_map = {} @@ -704,7 +713,7 @@ async def _save_background_health_checks_to_db( key = check.model_id if check.model_id else check.model_name if key not in latest_checks_map: latest_checks_map[key] = check - + # Step 4: Save aggregated results, but only if status changed await _save_health_check_results_if_changed( prisma_client, @@ -729,10 +738,15 @@ async def _perform_health_check_and_save( start_time, user_id, model_id=None, + max_concurrency=None, ): """Helper function to perform health check and save results to database""" healthy_endpoints, unhealthy_endpoints = await perform_health_check( - model_list=model_list, cli_model=cli_model, model=target_model, details=details + model_list=model_list, + cli_model=cli_model, + model=target_model, + details=details, + max_concurrency=max_concurrency, ) # Optionally save health check result to database (non-blocking) @@ -789,6 +803,7 @@ async def health_endpoint( import time from litellm.proxy.proxy_server import ( + health_check_concurrency, health_check_details, health_check_results, llm_model_list, @@ -841,6 +856,7 @@ async def health_endpoint( start_time=start_time, user_id=user_api_key_dict.user_id, model_id=None, # CLI model doesn't have model_id + max_concurrency=health_check_concurrency, ) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -864,6 +880,7 @@ async def health_endpoint( start_time=start_time, user_id=user_api_key_dict.user_id, model_id=model_id, + max_concurrency=health_check_concurrency, ) except Exception as e: verbose_proxy_logger.error( @@ -1420,11 +1437,11 @@ async def test_model_connection( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + # Get model name from litellm_params request_litellm_params = litellm_params or {} model_name = request_litellm_params.get("model") - + # Look up model configuration from router if model name is provided # This gets the litellm_params from proxy config (with resolved env vars) config_litellm_params: dict = {} @@ -1432,34 +1449,39 @@ async def test_model_connection( try: # First try to find by proxy model_name (e.g., "gpt-4o") deployments = llm_router.get_model_list(model_name=model_name) - + # If not found, try to find by litellm model name (e.g., "azure/gpt-4o") if not deployments or len(deployments) == 0: all_deployments = llm_router.get_model_list(model_name=None) if all_deployments: for deployment in all_deployments: - if deployment.get("litellm_params", {}).get("model") == model_name: + if ( + deployment.get("litellm_params", {}).get("model") + == model_name + ): deployments = [deployment] break - + if deployments and len(deployments) > 0: # Use the first deployment's litellm_params as base config # These already have resolved environment variables from proxy config - config_litellm_params = dict(deployments[0].get("litellm_params", {})) + config_litellm_params = dict( + deployments[0].get("litellm_params", {}) + ) except Exception as e: verbose_proxy_logger.debug( f"Could not find model {model_name} in router: {e}. " "Proceeding with request params only." ) - + # Merge: config params (from proxy config) as base, request params override # This allows users to override specific params while using config for credentials merged_litellm_params = {**config_litellm_params, **request_litellm_params} - + # Resolve os.environ/ environment variables in any remaining request params # This handles cases where user explicitly passes os.environ/ values to override config litellm_params = _resolve_os_environ_variables(merged_litellm_params) - + ## Auth check await ModelManagementAuthChecks.can_user_make_model_call( model_params=Deployment( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 35a75d4caf9..ddab560aaa3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9,6 +9,7 @@ import secrets import shutil import subprocess import sys +import threading import time import traceback import warnings @@ -658,7 +659,7 @@ _description = ( def cleanup_router_config_variables(): - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, prisma_client + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, health_check_concurrency, prisma_client # Set all variables to None master_key = None @@ -672,6 +673,7 @@ def cleanup_router_config_variables(): use_background_health_checks = None use_shared_health_check = None health_check_interval = None + health_check_concurrency = None prisma_client = None @@ -822,7 +824,9 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 verbose_proxy_logger.debug("About to initialize semantic tool filter") _config = proxy_config.get_config_state() _litellm_settings = _config.get("litellm_settings", {}) - verbose_proxy_logger.debug(f"litellm_settings keys = {list(_litellm_settings.keys())}") + verbose_proxy_logger.debug( + f"litellm_settings keys = {list(_litellm_settings.keys())}" + ) await ProxyStartupEvent._initialize_semantic_tool_filter( llm_router=llm_router, litellm_settings=_litellm_settings, @@ -1468,7 +1472,9 @@ redis_usage_cache: Optional[ RedisCache ] = None # redis cache used for tracking spend, tpm/rpm limits polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False -native_background_mode: List[str] = [] # Models that should use native provider background mode instead of polling +native_background_mode: List[ + str +] = [] # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None user_custom_key_generate = None @@ -1478,8 +1484,11 @@ use_background_health_checks = None use_shared_health_check = None use_queue = False health_check_interval = None +health_check_concurrency = None health_check_details = None health_check_results: Dict[str, Union[int, List[Dict[str, Any]]]] = {} +background_health_check_loop_active = False +background_health_check_cycle_seq = 0 queue: List = [] litellm_proxy_budget_name = "litellm-proxy-budget" litellm_proxy_admin_name = LITELLM_PROXY_ADMIN_NAME @@ -1927,6 +1936,29 @@ def run_ollama_serve(): ) +def _get_process_rss_mb() -> Optional[float]: + """ + Get process RSS memory in MB. + On Linux, ru_maxrss is in KB. On macOS, ru_maxrss is in bytes. + """ + try: + import resource + + ru_maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if sys.platform == "darwin": + return float(ru_maxrss) / (1024 * 1024) + return float(ru_maxrss) / 1024 + except Exception: + return None + + +def _rss_mb_for_log() -> str: + rss_mb = _get_process_rss_mb() + if rss_mb is None: + return "unknown" + return f"{rss_mb:.2f}" + + async def _run_background_health_check(): """ Periodically run health checks in the background on the endpoints. @@ -1934,7 +1966,10 @@ async def _run_background_health_check(): Update health_check_results, based on this. Uses shared health check state when Redis is available to coordinate across pods. """ - global health_check_results, llm_model_list, health_check_interval, health_check_details, use_shared_health_check, redis_usage_cache, prisma_client + global health_check_results, llm_model_list, health_check_interval + global health_check_concurrency, health_check_details, use_shared_health_check + global redis_usage_cache, prisma_client + global background_health_check_loop_active, background_health_check_cycle_seq if ( health_check_interval is None @@ -1943,6 +1978,24 @@ async def _run_background_health_check(): ): return + if background_health_check_loop_active: + verbose_proxy_logger.warning( + "background_health_check_loop_overlap_detected existing_loop_active=true interval_seconds=%s max_concurrency=%s shared=%s", + health_check_interval, + health_check_concurrency, + use_shared_health_check, + ) + background_health_check_loop_active = True + verbose_proxy_logger.info( + "background_health_check_loop_started interval_seconds=%s max_concurrency=%s shared=%s details=%s thread_count=%d rss_mb=%s", + health_check_interval, + health_check_concurrency, + use_shared_health_check, + health_check_details, + threading.active_count(), + _rss_mb_for_log(), + ) + # Initialize shared health check manager if Redis is available and feature is enabled shared_health_manager = None if use_shared_health_check and redis_usage_cache is not None: @@ -1958,8 +2011,13 @@ async def _run_background_health_check(): verbose_proxy_logger.info("Initialized shared health check manager") while True: + background_health_check_cycle_seq += 1 + cycle_id = f"bg-{background_health_check_cycle_seq}" + cycle_start_time = time.monotonic() + # make 1 deep copy of llm_model_list on every health check iteration _llm_model_list = copy.deepcopy(llm_model_list) or [] + model_count_total = len(_llm_model_list) # filter out models that have disabled background health checks _llm_model_list = [ @@ -1967,6 +2025,52 @@ async def _run_background_health_check(): for m in _llm_model_list if not m.get("model_info", {}).get("disable_background_health_check", False) ] + model_count_enabled = len(_llm_model_list) + expected_peak_in_flight = model_count_enabled + if ( + isinstance(health_check_concurrency, int) + and health_check_concurrency > 0 + and model_count_enabled > 0 + ): + expected_peak_in_flight = min(model_count_enabled, health_check_concurrency) + + verbose_proxy_logger.debug( + "background_health_check_cycle_start cycle_id=%s model_count_total=%d model_count_enabled=%d interval_seconds=%s max_concurrency=%s expected_peak_in_flight=%d shared=%s thread_count=%d rss_mb=%s", + cycle_id, + model_count_total, + model_count_enabled, + health_check_interval, + health_check_concurrency, + expected_peak_in_flight, + shared_health_manager is not None, + threading.active_count(), + _rss_mb_for_log(), + ) + + instrumentation_context = { + "enabled": True, + "source": "proxy_background_loop", + "cycle_id": cycle_id, + } + + async def _run_direct_health_check_with_instrumentation(): + try: + return await perform_health_check( + model_list=_llm_model_list, + details=health_check_details, + max_concurrency=health_check_concurrency, + instrumentation_context=instrumentation_context, + ) + except TypeError as e: + if "instrumentation_context" not in str(e): + raise + # Backward compatibility for monkeypatched or wrapped callables + # that do not accept instrumentation_context. + return await perform_health_check( + model_list=_llm_model_list, + details=health_check_details, + max_concurrency=health_check_concurrency, + ) # Use shared health check if available, otherwise fall back to direct health check # Convert health_check_details to bool for perform_shared_health_check (defaults to True if None) @@ -1980,19 +2084,21 @@ async def _run_background_health_check(): healthy_endpoints, unhealthy_endpoints, ) = await shared_health_manager.perform_shared_health_check( - model_list=_llm_model_list, details=details_bool + model_list=_llm_model_list, + details=details_bool, + max_concurrency=health_check_concurrency, ) except Exception as e: verbose_proxy_logger.error( "Error in shared health check, falling back to direct health check: %s", str(e), ) - healthy_endpoints, unhealthy_endpoints = await perform_health_check( - model_list=_llm_model_list, details=health_check_details + healthy_endpoints, unhealthy_endpoints = ( + await _run_direct_health_check_with_instrumentation() ) else: - healthy_endpoints, unhealthy_endpoints = await perform_health_check( - model_list=_llm_model_list, details=health_check_details + healthy_endpoints, unhealthy_endpoints = ( + await _run_direct_health_check_with_instrumentation() ) # Update the global variable with the health check results @@ -2000,6 +2106,25 @@ async def _run_background_health_check(): health_check_results["unhealthy_endpoints"] = unhealthy_endpoints health_check_results["healthy_count"] = len(healthy_endpoints) health_check_results["unhealthy_count"] = len(unhealthy_endpoints) + cycle_duration_ms = (time.monotonic() - cycle_start_time) * 1000 + verbose_proxy_logger.debug( + "background_health_check_cycle_complete cycle_id=%s model_count_enabled=%d healthy_count=%d unhealthy_count=%d duration_ms=%.2f interval_seconds=%s thread_count=%d rss_mb=%s", + cycle_id, + model_count_enabled, + len(healthy_endpoints), + len(unhealthy_endpoints), + cycle_duration_ms, + health_check_interval, + threading.active_count(), + _rss_mb_for_log(), + ) + if cycle_duration_ms > (health_check_interval * 1000): + verbose_proxy_logger.warning( + "background_health_check_cycle_duration_exceeded_interval cycle_id=%s duration_ms=%.2f interval_seconds=%s", + cycle_id, + cycle_duration_ms, + health_check_interval, + ) # Save background health checks to database (non-blocking) if prisma_client is not None: @@ -2480,7 +2605,7 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, proxy_batch_polling_interval, config_passthrough_endpoints + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, health_check_concurrency, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, proxy_batch_polling_interval, config_passthrough_endpoints config: dict = await self.get_config(config_file_path=config_file_path) @@ -2905,7 +3030,18 @@ class ProxyConfig: health_check_interval = general_settings.get( "health_check_interval", DEFAULT_HEALTH_CHECK_INTERVAL ) + health_check_concurrency = general_settings.get( + "health_check_concurrency", None + ) health_check_details = general_settings.get("health_check_details", True) + verbose_proxy_logger.info( + "background_health_check_config enabled=%s shared=%s interval_seconds=%s max_concurrency=%s details=%s", + use_background_health_checks, + use_shared_health_check, + health_check_interval, + health_check_concurrency, + health_check_details, + ) ### RBAC ### rbac_role_permissions = general_settings.get("role_permissions", None) @@ -2999,7 +3135,7 @@ class ProxyConfig: for k, v in router_settings.items(): if k in available_args: router_params[k] = v - elif k == "health_check_interval": + elif k in {"health_check_interval", "health_check_concurrency"}: raise ValueError( f"'{k}' is NOT a valid router_settings parameter. Please move it to 'general_settings'." ) @@ -4201,9 +4337,7 @@ class ProxyConfig: ) if self._should_load_db_object(object_type="semantic_filter_settings"): - await self._init_semantic_filter_settings_in_db( - prisma_client=prisma_client - ) + await self._init_semantic_filter_settings_in_db(prisma_client=prisma_client) async def _init_semantic_filter_settings_in_db(self, prisma_client: PrismaClient): """ @@ -5259,30 +5393,38 @@ class ProxyStartupEvent: ): """Initialize MCP semantic tool filter if configured""" from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook - - mcp_semantic_filter_config = litellm_settings.get("mcp_semantic_tool_filter", None) - + + mcp_semantic_filter_config = litellm_settings.get( + "mcp_semantic_tool_filter", None + ) + # Only proceed if the feature is configured and enabled - if not mcp_semantic_filter_config or not mcp_semantic_filter_config.get("enabled", False): - verbose_proxy_logger.debug("Semantic tool filter not configured or not enabled, skipping initialization") + if not mcp_semantic_filter_config or not mcp_semantic_filter_config.get( + "enabled", False + ): + verbose_proxy_logger.debug( + "Semantic tool filter not configured or not enabled, " + "skipping initialization" + ) return - + verbose_proxy_logger.debug( f"Initializing semantic tool filter: llm_router={llm_router is not None}, " f"config={mcp_semantic_filter_config}" ) - hook = await SemanticToolFilterHook.initialize_from_config( config=mcp_semantic_filter_config, llm_router=llm_router, ) - + if hook: verbose_proxy_logger.debug("Semantic tool filter hook registered") litellm.logging_callback_manager.add_litellm_callback(hook) else: # Only warn if the feature was configured but failed to initialize - verbose_proxy_logger.warning("Semantic tool filter hook was configured but failed to initialize") + verbose_proxy_logger.warning( + "Semantic tool filter hook was configured but failed to initialize" + ) @classmethod def _initialize_jwt_auth( @@ -8706,7 +8848,8 @@ async def _apply_search_filter_to_models( # Fetch database models if we need more for the current page if router_models_count < models_needed_for_page: models_to_fetch = min( - models_needed_for_page - router_models_count, db_models_total_count + models_needed_for_page - router_models_count, + db_models_total_count, ) if models_to_fetch > 0: @@ -8742,21 +8885,21 @@ async def _apply_search_filter_to_models( def _normalize_datetime_for_sorting(dt: Any) -> Optional[datetime]: """ Normalize a datetime value to a timezone-aware UTC datetime for sorting. - + This function handles: - None values: returns None - String values: parses ISO format strings and converts to UTC-aware datetime - Datetime objects: converts naive datetimes to UTC-aware, and aware datetimes to UTC - + Args: dt: Datetime value (None, str, or datetime object) - + Returns: UTC-aware datetime object, or None if input is None or cannot be parsed """ if dt is None: return None - + if isinstance(dt, str): try: # Handle ISO format strings, including 'Z' suffix @@ -8770,14 +8913,14 @@ def _normalize_datetime_for_sorting(dt: Any) -> Optional[datetime]: return parsed_dt except (ValueError, AttributeError): return None - + if isinstance(dt, datetime): # If naive, assume UTC and make it aware if dt.tzinfo is None: return dt.replace(tzinfo=timezone.utc) # If aware, convert to UTC return dt.astimezone(timezone.utc) - + return None @@ -8797,46 +8940,60 @@ def _sort_models( Returns: Sorted list of models """ - if not sort_by or sort_by not in ["model_name", "created_at", "updated_at", "costs", "status"]: + if not sort_by or sort_by not in [ + "model_name", + "created_at", + "updated_at", + "costs", + "status", + ]: return all_models reverse = sort_order.lower() == "desc" def get_sort_key(model: Dict[str, Any]) -> Any: model_info = model.get("model_info", {}) - + if sort_by == "model_name": return model.get("model_name", "").lower() - + elif sort_by == "created_at": created_at = model_info.get("created_at") normalized_dt = _normalize_datetime_for_sorting(created_at) if normalized_dt is None: # Put None values at the end for asc, at the start for desc - return (datetime.max.replace(tzinfo=timezone.utc) if not reverse else datetime.min.replace(tzinfo=timezone.utc)) + return ( + datetime.max.replace(tzinfo=timezone.utc) + if not reverse + else datetime.min.replace(tzinfo=timezone.utc) + ) return normalized_dt - + elif sort_by == "updated_at": updated_at = model_info.get("updated_at") normalized_dt = _normalize_datetime_for_sorting(updated_at) if normalized_dt is None: - return (datetime.max.replace(tzinfo=timezone.utc) if not reverse else datetime.min.replace(tzinfo=timezone.utc)) + return ( + datetime.max.replace(tzinfo=timezone.utc) + if not reverse + else datetime.min.replace(tzinfo=timezone.utc) + ) return normalized_dt - + elif sort_by == "costs": input_cost = model_info.get("input_cost_per_token", 0) or 0 output_cost = model_info.get("output_cost_per_token", 0) or 0 total_cost = input_cost + output_cost # Put 0 or None costs at the end for asc, at the start for desc if total_cost == 0: - return (float("inf") if not reverse else float("-inf")) + return float("inf") if not reverse else float("-inf") return total_cost - + elif sort_by == "status": # False (config) comes before True (db) for asc db_model = model_info.get("db_model", False) return db_model - + return None try: @@ -9032,9 +9189,7 @@ async def _find_model_by_id( ) if db_model: # Convert database model to router format - decrypted_models = proxy_config.decrypt_model_list_from_db( - [db_model] - ) + decrypted_models = proxy_config.decrypt_model_list_from_db([db_model]) if decrypted_models: found_model = decrypted_models[0] except Exception as e: @@ -9208,13 +9363,13 @@ async def model_info_v2( ) verbose_proxy_logger.debug("all_models: %s", all_models) - + # Append A2A agents to models list all_models = await append_agents_to_model_info( models=all_models, user_api_key_dict=user_api_key_dict, ) - + # Update total count to include agents search_total_count = len(all_models) @@ -10057,7 +10212,7 @@ async def model_group_info( model_groups: List[ModelGroupInfoProxy] = _get_model_group_info( llm_router=llm_router, all_models_str=all_models_str, model_group=model_group ) - + # Append A2A agents to model groups model_groups = await append_agents_to_model_group( model_groups=model_groups, diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index 19882bbe4be..7ea4574bab9 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -92,7 +92,7 @@ async def test_azure_img_gen_health_check(): litellm._turn_on_debug() max_retries = 3 retry_delay = 1 # Start with 1 second delay - + for attempt in range(max_retries): response = await litellm.ahealth_check( model_params={ @@ -103,11 +103,11 @@ async def test_azure_img_gen_health_check(): mode="image_generation", prompt="cute baby sea otter", ) - + # Check if response is successful (no error) if isinstance(response, dict) and "error" not in response: return response - + # Check if error is a transient Azure internal server error error_str = str(response.get("error", "")).lower() is_transient_error = ( @@ -116,16 +116,18 @@ async def test_azure_img_gen_health_check(): or "internalfailure" in error_str or "internal failure" in error_str ) - + # If it's the last attempt or not a transient error, fail the test if attempt == max_retries - 1 or not is_transient_error: - assert isinstance(response, dict) and "error" not in response, f"Health check failed: {response.get('error', 'Unknown error')}" + assert ( + isinstance(response, dict) and "error" not in response + ), f"Health check failed: {response.get('error', 'Unknown error')}" return response - + # Wait before retrying with exponential backoff await asyncio.sleep(retry_delay) retry_delay *= 2 # Exponential backoff - + # Should not reach here, but just in case assert False, "Health check failed after all retries" @@ -562,6 +564,99 @@ async def test_health_check_bad_model(): ), "Health check took longer than health_check_timeout" +@pytest.mark.asyncio +async def test_health_check_respects_concurrency_limit(): + from litellm.proxy.health_check import _perform_health_check + + model_list = [ + {"litellm_params": {"model": f"openai/gpt-4o-mini-{i}", "api_key": "fake-key"}} + for i in range(6) + ] + + active = 0 + max_active = 0 + + async def mock_health_check(litellm_params, **kwargs): + nonlocal active, max_active + active += 1 + max_active = max(max_active, active) + await asyncio.sleep(0.05) + active -= 1 + return {"status": "healthy"} + + with patch("litellm.ahealth_check", side_effect=mock_health_check): + await _perform_health_check(model_list, max_concurrency=2) + + assert max_active <= 2 + + +@pytest.mark.asyncio +async def test_health_check_creates_only_bounded_initial_tasks(): + from litellm.proxy.health_check import _perform_health_check + + model_list = [ + {"litellm_params": {"model": f"openai/gpt-4o-mini-{i}", "api_key": "fake-key"}} + for i in range(10) + ] + release_event = asyncio.Event() + create_task_call_count = 0 + real_create_task = asyncio.create_task + + async def mock_health_check(litellm_params, **kwargs): + await release_event.wait() + return {"status": "healthy"} + + def tracked_create_task(coro): + nonlocal create_task_call_count + create_task_call_count += 1 + return real_create_task(coro) + + with patch("litellm.ahealth_check", side_effect=mock_health_check), patch( + "litellm.proxy.health_check.asyncio.create_task", side_effect=tracked_create_task + ): + perform_task = real_create_task( + _perform_health_check(model_list, max_concurrency=2) + ) + await asyncio.sleep(0.05) + assert create_task_call_count == 2 + release_event.set() + await perform_task + + +@pytest.mark.asyncio +async def test_timeout_does_not_cancel_other_health_checks(): + from litellm.proxy.health_check import _perform_health_check + + model_list = [ + { + "litellm_params": {"model": "openai/slow-model", "api_key": "fake-key"}, + "model_info": {"health_check_timeout": 0.05}, + }, + { + "litellm_params": {"model": "openai/fast-model", "api_key": "fake-key"}, + "model_info": {"health_check_timeout": 1}, + }, + ] + + async def mock_health_check(litellm_params, **kwargs): + if litellm_params["model"] == "openai/slow-model": + await asyncio.sleep(0.2) + return {"status": "healthy"} + await asyncio.sleep(0.01) + return {"status": "healthy"} + + with patch("litellm.ahealth_check", side_effect=mock_health_check): + healthy_endpoints, unhealthy_endpoints = await _perform_health_check( + model_list, max_concurrency=1 + ) + + healthy_models = {endpoint["model"] for endpoint in healthy_endpoints} + unhealthy_models = {endpoint["model"] for endpoint in unhealthy_endpoints} + + assert "openai/fast-model" in healthy_models + assert "openai/slow-model" in unhealthy_models + + @pytest.mark.asyncio async def test_ahealth_check_ocr(): litellm._turn_on_debug() @@ -643,20 +738,20 @@ async def test_image_generation_health_check_prompt(monkeypatch): async def test_health_check_with_custom_llm_provider(): """ Test that ahealth_check correctly uses custom_llm_provider from model_params. - + This test verifies the fix for the issue where the UI's "Test connect" button failed with "LLM Provider NOT provided" error for OpenAI-compatible self-hosted providers, even when a provider was selected in the dropdown. - + The fix ensures that when custom_llm_provider is passed in model_params, it's properly forwarded to get_llm_provider() to identify the correct provider. """ from unittest.mock import MagicMock - + # Mock the completion call to avoid making real API calls mock_response = MagicMock() mock_response._hidden_params = {"headers": {"x-ratelimit-remaining-tokens": "1000"}} - + with patch("litellm.acompletion", return_value=mock_response): # Test with a custom model name that wouldn't be recognized without custom_llm_provider response = await litellm.ahealth_check( @@ -668,7 +763,7 @@ async def test_health_check_with_custom_llm_provider(): }, mode="chat", ) - + # Should succeed without "LLM Provider NOT provided" error assert "error" not in response assert isinstance(response, dict) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 14b49901c9c..0cbba7b5cc3 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2330,7 +2330,9 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch): test_model_list_2 = [{"model_name": "model-b"}] called_model_lists = [] - async def fake_perform_health_check(model_list, details): + async def fake_perform_health_check( + model_list, details, max_concurrency=None + ): called_model_lists.append(copy.deepcopy(model_list)) return (["healthy"], ["unhealthy"]) @@ -2378,7 +2380,9 @@ async def test_background_health_check_skip_disabled_models(monkeypatch): ] called_model_lists = [] - async def fake_perform_health_check(model_list, details): + async def fake_perform_health_check( + model_list, details, max_concurrency=None + ): called_model_lists.append(copy.deepcopy(model_list)) return (["healthy"], []) From 5e9f24f74c59d0272b050a17d0b34ae82646e9fb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 09:32:11 -0800 Subject: [PATCH 017/132] fix(bedrock): pass timeout param to bedrock rerank http client (#22021) * fix(bedrock): pass timeout to bedrock rerank http client * refactor: extract large functions to fix PLR0915 ruff lint errors --- .../litellm_core_utils/realtime_streaming.py | 248 +++++++++--------- litellm/llms/bedrock/rerank/handler.py | 8 +- ...odel_prices_and_context_window_backup.json | 95 +++++-- litellm/proxy/health_check.py | 122 ++++----- litellm/proxy/proxy_server.py | 127 +++++---- litellm/rerank_api/main.py | 1 + .../test_bedrock_rerank_header_forwarding.py | 77 ++++++ 7 files changed, 427 insertions(+), 251 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index b42bf9e5711..759eaf60035 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -203,6 +203,129 @@ class RealTimeStreaming: return True return False + async def _handle_provider_config_message(self, raw_response) -> None: + """Process a backend message when a provider_config is set (transformed path).""" + returned_object = self.provider_config.transform_realtime_response( # type: ignore[union-attr] + raw_response, + self.model, + self.logging_obj, + realtime_response_transform_input={ + "session_configuration_request": self.session_configuration_request, + "current_output_item_id": self.current_output_item_id, + "current_response_id": self.current_response_id, + "current_delta_chunks": self.current_delta_chunks, + "current_conversation_id": self.current_conversation_id, + "current_item_chunks": self.current_item_chunks, + "current_delta_type": self.current_delta_type, + }, + ) + + transformed_response = returned_object["response"] + self.current_output_item_id = returned_object["current_output_item_id"] + self.current_response_id = returned_object["current_response_id"] + self.current_delta_chunks = returned_object["current_delta_chunks"] + self.current_conversation_id = returned_object["current_conversation_id"] + self.current_item_chunks = returned_object["current_item_chunks"] + self.current_delta_type = returned_object["current_delta_type"] + self.session_configuration_request = returned_object["session_configuration_request"] + events = ( + transformed_response + if isinstance(transformed_response, list) + else [transformed_response] + ) + for event in events: + ## GUARDRAIL: inject create_response=false on session.created + if isinstance(event, dict) and event.get("type") == "session.created": + if self._has_realtime_guardrails(): + await self.backend_ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "turn_detection": { + "type": "server_vad", + "create_response": False, + } + }, + } + ) + ) + for event in events: + event_str = json.dumps(event) + ## GUARDRAIL: run on transcription events in provider_config path too + if ( + isinstance(event, dict) + and event.get("type") + == "conversation.item.input_audio_transcription.completed" + ): + transcript = event.get("transcript", "") + self.store_message(event_str) + await self.websocket.send_text(event_str) + blocked = await self.run_realtime_guardrails( + transcript, item_id=event.get("item_id") + ) + if not blocked: + await self.backend_ws.send( + json.dumps({"type": "response.create"}) + ) + continue + ## LOGGING + self.store_message(event_str) + await self.websocket.send_text(event_str) + + async def _handle_raw_backend_message(self, raw_response) -> bool: + """Process a backend message without provider_config (raw path). + + Returns True if the caller should skip the default store+forward (i.e. continue the loop). + """ + try: + event_obj = json.loads(raw_response) + + if event_obj.get("type") == "session.created": + # If any realtime guardrails are registered, proactively + # set create_response=false so the LLM never auto-responds + # before our guardrail has a chance to run. + if self._has_realtime_guardrails(): + await self.backend_ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "turn_detection": { + "type": "server_vad", + "create_response": False, + } + }, + } + ) + ) + verbose_logger.debug( + "[realtime guardrail] injected create_response=false into session" + ) + + if ( + event_obj.get("type") + == "conversation.item.input_audio_transcription.completed" + ): + transcript = event_obj.get("transcript", "") + ## LOGGING — must happen before continue below + self.store_message(raw_response) + # Forward transcript to client so user sees what they said + await self.websocket.send_text(raw_response) + blocked = await self.run_realtime_guardrails( + transcript, + item_id=event_obj.get("item_id"), + ) + if not blocked: + # Clean — trigger LLM response + await self.backend_ws.send( + json.dumps({"type": "response.create"}) + ) + return True + except (json.JSONDecodeError, AttributeError): + pass + return False + async def backend_to_client_send_messages(self): import websockets @@ -216,128 +339,11 @@ class RealTimeStreaming: raw_response = await self.backend_ws.recv() # type: ignore[assignment] if self.provider_config: - returned_object = self.provider_config.transform_realtime_response( - raw_response, - self.model, - self.logging_obj, - realtime_response_transform_input={ - "session_configuration_request": self.session_configuration_request, - "current_output_item_id": self.current_output_item_id, - "current_response_id": self.current_response_id, - "current_delta_chunks": self.current_delta_chunks, - "current_conversation_id": self.current_conversation_id, - "current_item_chunks": self.current_item_chunks, - "current_delta_type": self.current_delta_type, - }, - ) - - transformed_response = returned_object["response"] - self.current_output_item_id = returned_object[ - "current_output_item_id" - ] - self.current_response_id = returned_object["current_response_id"] - self.current_delta_chunks = returned_object["current_delta_chunks"] - self.current_conversation_id = returned_object[ - "current_conversation_id" - ] - self.current_item_chunks = returned_object["current_item_chunks"] - self.current_delta_type = returned_object["current_delta_type"] - self.session_configuration_request = returned_object[ - "session_configuration_request" - ] - events = ( - transformed_response - if isinstance(transformed_response, list) - else [transformed_response] - ) - for event in events: - ## GUARDRAIL: inject create_response=false on session.created - if isinstance(event, dict) and event.get("type") == "session.created": - if self._has_realtime_guardrails(): - await self.backend_ws.send( - json.dumps( - { - "type": "session.update", - "session": { - "turn_detection": { - "type": "server_vad", - "create_response": False, - } - }, - } - ) - ) - for event in events: - event_str = json.dumps(event) - ## GUARDRAIL: run on transcription events in provider_config path too - if ( - isinstance(event, dict) - and event.get("type") - == "conversation.item.input_audio_transcription.completed" - ): - transcript = event.get("transcript", "") - self.store_message(event_str) - await self.websocket.send_text(event_str) - blocked = await self.run_realtime_guardrails( - transcript, item_id=event.get("item_id") - ) - if not blocked: - await self.backend_ws.send( - json.dumps({"type": "response.create"}) - ) - continue - ## LOGGING - self.store_message(event_str) - await self.websocket.send_text(event_str) - + await self._handle_provider_config_message(raw_response) else: - ## GUARDRAIL: intercept transcription events before triggering LLM - try: - event_obj = json.loads(raw_response) - - if event_obj.get("type") == "session.created": - # If any realtime guardrails are registered, proactively - # set create_response=false so the LLM never auto-responds - # before our guardrail has a chance to run. - if self._has_realtime_guardrails(): - await self.backend_ws.send( - json.dumps( - { - "type": "session.update", - "session": { - "turn_detection": { - "type": "server_vad", - "create_response": False, - } - }, - } - ) - ) - verbose_logger.debug( - "[realtime guardrail] injected create_response=false into session" - ) - - if ( - event_obj.get("type") - == "conversation.item.input_audio_transcription.completed" - ): - transcript = event_obj.get("transcript", "") - ## LOGGING — must happen before continue below - self.store_message(raw_response) - # Forward transcript to client so user sees what they said - await self.websocket.send_text(raw_response) - blocked = await self.run_realtime_guardrails( - transcript, - item_id=event_obj.get("item_id"), - ) - if not blocked: - # Clean — trigger LLM response - await self.backend_ws.send( - json.dumps({"type": "response.create"}) - ) - continue - except (json.JSONDecodeError, AttributeError): - pass + handled = await self._handle_raw_backend_message(raw_response) + if handled: + continue ## LOGGING self.store_message(raw_response) await self.websocket.send_text(raw_response) diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 06f1e9e86c9..37167e7c330 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -29,12 +29,13 @@ class BedrockRerankHandler(BaseAWSLLM): async def arerank( self, prepared_request: BedrockPreparedRequest, + timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[AsyncHTTPHandler] = None, ): if client is None: client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) try: - response = await client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"]) + response = await client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"], timeout=timeout) response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code @@ -56,6 +57,7 @@ class BedrockRerankHandler(BaseAWSLLM): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, _is_async: Optional[bool] = False, + timeout: Optional[Union[float, httpx.Timeout]] = None, api_base: Optional[str] = None, extra_headers: Optional[dict] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, @@ -89,12 +91,12 @@ class BedrockRerankHandler(BaseAWSLLM): ) if _is_async: - return self.arerank(prepared_request, client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None) # type: ignore + return self.arerank(prepared_request, timeout=timeout, client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None) # type: ignore if client is None or not isinstance(client, HTTPHandler): client = _get_httpx_client() try: - response = client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"]) + response = client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"], timeout=timeout) response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3909ce4c8b0..624714feb87 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26555,65 +26555,124 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "perplexity/preset/fast-search": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_preset": true, + "supports_function_calling": true + }, "perplexity/preset/pro-search": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_preset": true + "supports_preset": true, + "supports_function_calling": true }, - "perplexity/openai/gpt-4o": { + "perplexity/preset/deep-research": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_preset": true, + "supports_function_calling": true }, - "perplexity/openai/gpt-4o-mini": { + "perplexity/preset/advanced-deep-research": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_preset": true, + "supports_function_calling": true }, "perplexity/openai/gpt-5.2": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": true + "supports_reasoning": true, + "supports_function_calling": true }, - "perplexity/anthropic/claude-3-5-sonnet-20241022": { + "perplexity/openai/gpt-5.1": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/anthropic/claude-3-5-haiku-20241022": { + "perplexity/openai/gpt-5-mini": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/google/gemini-2.0-flash-exp": { + "perplexity/anthropic/claude-opus-4-6": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/google/gemini-2.0-flash-thinking-exp": { + "perplexity/anthropic/claude-opus-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": true + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/xai/grok-2-1212": { + "perplexity/anthropic/claude-sonnet-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/xai/grok-2-vision-1212": { + "perplexity/anthropic/claude-haiku-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-3-pro-preview": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-3-flash-preview": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-2.5-pro": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-2.5-flash": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/xai/grok-4-1-fast-non-reasoning": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/perplexity/sonar": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true }, "publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": { "input_cost_per_token": 0.0, diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 4777f64405d..de4517efc88 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -11,7 +11,7 @@ from typing import List, Optional import litellm logger = logging.getLogger(__name__) -from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS, DEFAULT_HEALTH_CHECK_PROMPT +from litellm.constants import DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS ILLEGAL_DISPLAY_PARAMS = [ "messages", @@ -98,6 +98,66 @@ async def run_with_timeout(task, timeout): return {"error": "Timeout exceeded"} +async def _run_model_health_check(model: dict): + litellm_params = model["litellm_params"] + model_info = model.get("model_info", {}) + mode = model_info.get("mode", None) + litellm_params = _update_litellm_params_for_health_check(model_info, litellm_params) + timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS + + return await run_with_timeout( + litellm.ahealth_check( + litellm_params, + mode=mode, + prompt=DEFAULT_HEALTH_CHECK_PROMPT, + input=["test from litellm"], + ), + timeout, + ) + + +async def _run_health_checks_with_bounded_concurrency( + models: list, concurrency_limit: int +) -> tuple[list, int]: + """ + Run health checks with at most `concurrency_limit` active tasks. + Preserves result ordering to match `models`. + """ + results: list = [None] * len(models) + tasks_to_index: dict[asyncio.Task, int] = {} + model_iter = iter(enumerate(models)) + peak_in_flight = 0 + + def _schedule_next() -> bool: + nonlocal peak_in_flight + try: + idx, next_model = next(model_iter) + except StopIteration: + return False + task = asyncio.create_task(_run_model_health_check(next_model)) + tasks_to_index[task] = idx + peak_in_flight = max(peak_in_flight, len(tasks_to_index)) + return True + + for _ in range(min(concurrency_limit, len(models))): + _schedule_next() + + while tasks_to_index: + done, _ = await asyncio.wait( + set(tasks_to_index.keys()), + return_when=asyncio.FIRST_COMPLETED, + ) + for task in done: + idx = tasks_to_index.pop(task) + try: + results[idx] = task.result() + except Exception as e: + results[idx] = e + _schedule_next() + + return results, peak_in_flight + + async def _perform_health_check( model_list: list, details: Optional[bool] = True, @@ -115,66 +175,6 @@ async def _perform_health_check( cycle_id = instrumentation_context.get("cycle_id", "unknown") source = instrumentation_context.get("source", "unknown") - async def _run_model_health_check(model: dict): - litellm_params = model["litellm_params"] - model_info = model.get("model_info", {}) - mode = model_info.get("mode", None) - litellm_params = _update_litellm_params_for_health_check( - model_info, litellm_params - ) - timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS - - return await run_with_timeout( - litellm.ahealth_check( - litellm_params, - mode=mode, - prompt=DEFAULT_HEALTH_CHECK_PROMPT, - input=["test from litellm"], - ), - timeout, - ) - - async def _run_health_checks_with_bounded_concurrency( - models: list, concurrency_limit: int - ) -> tuple[list, int]: - """ - Run health checks with at most `concurrency_limit` active tasks. - Preserves result ordering to match `models`. - """ - results: list = [None] * len(models) - tasks_to_index: dict[asyncio.Task, int] = {} - model_iter = iter(enumerate(models)) - peak_in_flight = 0 - - def _schedule_next() -> bool: - nonlocal peak_in_flight - try: - idx, next_model = next(model_iter) - except StopIteration: - return False - task = asyncio.create_task(_run_model_health_check(next_model)) - tasks_to_index[task] = idx - peak_in_flight = max(peak_in_flight, len(tasks_to_index)) - return True - - for _ in range(min(concurrency_limit, len(models))): - _schedule_next() - - while tasks_to_index: - done, _ = await asyncio.wait( - set(tasks_to_index.keys()), - return_when=asyncio.FIRST_COMPLETED, - ) - for task in done: - idx = tasks_to_index.pop(task) - try: - results[idx] = task.result() - except Exception as e: - results[idx] = e - _schedule_next() - - return results, peak_in_flight - dispatch_mode = "unbounded" peak_in_flight = 0 if isinstance(max_concurrency, int) and max_concurrency > 0: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ddab560aaa3..b3b6bf0ccf7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -713,7 +713,7 @@ async def _initialize_shared_aiohttp_session(): try: from aiohttp import ClientSession, TCPConnector - connector_kwargs = { + connector_kwargs: Dict[str, Any] = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, } @@ -1959,6 +1959,65 @@ def _rss_mb_for_log() -> str: return f"{rss_mb:.2f}" +async def _run_direct_health_check_with_instrumentation( + model_list: list, + details: Optional[bool], + max_concurrency: Optional[int], + instrumentation_context: dict, +): + try: + return await perform_health_check( + model_list=model_list, + details=details, + max_concurrency=max_concurrency, + instrumentation_context=instrumentation_context, + ) + except TypeError as e: + if "instrumentation_context" not in str(e): + raise + # Backward compatibility for monkeypatched or wrapped callables + # that do not accept instrumentation_context. + return await perform_health_check( + model_list=model_list, + details=details, + max_concurrency=max_concurrency, + ) + + +def _schedule_background_health_check_db_save( + prisma_client, + shared_health_manager, + model_list: list, + healthy_endpoints: list, + unhealthy_endpoints: list, +): + """Fire-and-forget: persist health check results to DB if prisma is available.""" + if prisma_client is None: + return + import time as time_module + + from litellm.proxy.health_endpoints._health_endpoints import ( + _save_background_health_checks_to_db, + ) + + checked_by = ( + shared_health_manager.pod_id + if shared_health_manager is not None + else "background_health_check" + ) + start_time = time_module.time() + asyncio.create_task( + _save_background_health_checks_to_db( + prisma_client, + model_list, + healthy_endpoints, + unhealthy_endpoints, + start_time, + checked_by=checked_by, + ) + ) + + async def _run_background_health_check(): """ Periodically run health checks in the background on the endpoints. @@ -2053,25 +2112,6 @@ async def _run_background_health_check(): "cycle_id": cycle_id, } - async def _run_direct_health_check_with_instrumentation(): - try: - return await perform_health_check( - model_list=_llm_model_list, - details=health_check_details, - max_concurrency=health_check_concurrency, - instrumentation_context=instrumentation_context, - ) - except TypeError as e: - if "instrumentation_context" not in str(e): - raise - # Backward compatibility for monkeypatched or wrapped callables - # that do not accept instrumentation_context. - return await perform_health_check( - model_list=_llm_model_list, - details=health_check_details, - max_concurrency=health_check_concurrency, - ) - # Use shared health check if available, otherwise fall back to direct health check # Convert health_check_details to bool for perform_shared_health_check (defaults to True if None) details_bool = ( @@ -2094,11 +2134,21 @@ async def _run_background_health_check(): str(e), ) healthy_endpoints, unhealthy_endpoints = ( - await _run_direct_health_check_with_instrumentation() + await _run_direct_health_check_with_instrumentation( + _llm_model_list, + health_check_details, + health_check_concurrency, + instrumentation_context, + ) ) else: healthy_endpoints, unhealthy_endpoints = ( - await _run_direct_health_check_with_instrumentation() + await _run_direct_health_check_with_instrumentation( + _llm_model_list, + health_check_details, + health_check_concurrency, + instrumentation_context, + ) ) # Update the global variable with the health check results @@ -2127,32 +2177,13 @@ async def _run_background_health_check(): ) # Save background health checks to database (non-blocking) - if prisma_client is not None: - import time as time_module - - from litellm.proxy.health_endpoints._health_endpoints import ( - _save_background_health_checks_to_db, - ) - - # Use pod_id or a system identifier for checked_by if shared health check is enabled - checked_by = None - if shared_health_manager is not None: - checked_by = shared_health_manager.pod_id - else: - # Use a system identifier for background health checks - checked_by = "background_health_check" - - start_time = time_module.time() - asyncio.create_task( - _save_background_health_checks_to_db( - prisma_client, - _llm_model_list, - healthy_endpoints, - unhealthy_endpoints, - start_time, - checked_by=checked_by, - ) - ) + _schedule_background_health_check_db_save( + prisma_client, + shared_health_manager, + _llm_model_list, + healthy_endpoints, + unhealthy_endpoints, + ) await asyncio.sleep(health_check_interval) diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index f47fd6323f0..871b18062ff 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -365,6 +365,7 @@ def rerank( # noqa: PLR0915 max_chunks_per_doc=max_chunks_per_doc, _is_async=_is_async, optional_params=optional_params.model_dump(exclude_unset=True), + timeout=optional_params.timeout, api_base=api_base, extra_headers=merged_headers, logging_obj=litellm_logging_obj, diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index a8ac680908e..d6dc4bfa48d 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -237,6 +237,83 @@ async def test_bedrock_rerank_header_forwarding_async(model): pytest.fail(f"Failed to forward headers to {model}: {str(e)}") +def test_bedrock_rerank_timeout_sync(): + """ + Test that the timeout parameter is passed through to the HTTP client for Bedrock rerank (sync). + """ + client = HTTPHandler() + model = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0" + mock_credentials_info = create_mock_credentials() + + with patch.object(client, "post") as mock_post, \ + patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \ + patch("botocore.auth.SigV4Auth") as mock_sigv4: + + mock_sigv4.return_value = MagicMock() + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(bedrock_rerank_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_response.raise_for_status = lambda: None + mock_post.return_value = mock_response + + litellm.rerank( + model=model, + query=test_query, + documents=test_documents, + top_n=3, + client=client, + timeout=0.001, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + + assert mock_post.called + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs.get("timeout") == 0.001, ( + f"Expected timeout=0.001, got timeout={call_kwargs.get('timeout')}" + ) + + +@pytest.mark.asyncio +async def test_bedrock_rerank_timeout_async(): + """ + Test that the timeout parameter is passed through to the HTTP client for Bedrock rerank (async). + """ + client = AsyncHTTPHandler() + model = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0" + mock_credentials_info = create_mock_credentials() + + with patch.object(client, "post", new_callable=AsyncMock) as mock_post, \ + patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \ + patch("botocore.auth.SigV4Auth") as mock_sigv4: + + mock_sigv4.return_value = MagicMock() + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.text = json.dumps(bedrock_rerank_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_response.raise_for_status = lambda: None + mock_post.return_value = mock_response + + await litellm.arerank( + model=model, + query=test_query, + documents=test_documents, + top_n=3, + client=client, + timeout=0.001, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + + assert mock_post.called + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs.get("timeout") == 0.001, ( + f"Expected timeout=0.001, got timeout={call_kwargs.get('timeout')}" + ) + + def test_bedrock_rerank_extra_headers_and_headers_merge(): """ Test that both extra_headers and headers parameters are correctly merged for Bedrock rerank. From 86ec2fbf843c8e2bba0704f337e5b26423615bc4 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 24 Feb 2026 09:56:16 -0800 Subject: [PATCH 018/132] perf(proxy): batch 11 create_task() calls into 1 in update_database() Replace 11 separate asyncio.create_task() calls per request with a single batched task that runs all spend-update helpers sequentially. This reduces task scheduling overhead at high RPS (11,000 -> 1,000 tasks/sec at 1K RPS) and cuts 5 copy.deepcopy(payload) calls to 1 shared copy. Also fixes a mutation bug where the daily agent spend handler received the raw payload without deepcopy, unlike all other daily helpers. --- litellm/proxy/db/db_spend_update_writer.py | 284 +++++++++++------ .../proxy/db/test_db_spend_update_writer.py | 293 +++++++++++++----- 2 files changed, 419 insertions(+), 158 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 03628fda47f..d7d065d9064 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -124,53 +124,20 @@ class DBSpendUpdateWriter: payload["startTime"] = payload["startTime"].isoformat() if isinstance(payload["endTime"], datetime): payload["endTime"] = payload["endTime"].isoformat() - + if org_id is not None and org_id != "": payload["organization_id"] = org_id if team_id is not None and team_id != "": payload["team_id"] = team_id - asyncio.create_task( - self._update_user_db( - response_cost=response_cost, - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - litellm_proxy_budget_name=litellm_proxy_budget_name, - end_user_id=end_user_id, - ) - ) - asyncio.create_task( - self._update_key_db( - response_cost=response_cost, - hashed_token=hashed_token, - prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self._update_team_db( - response_cost=response_cost, - team_id=team_id, - user_id=user_id, - prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self._update_org_db( - response_cost=response_cost, - org_id=org_id, - prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self._update_tag_db( - response_cost=response_cost, - request_tags=copy.deepcopy(payload.get("request_tags")), - prisma_client=prisma_client, - ) - ) + # One deepcopy shared by all 6 daily spend helpers (was 5, fixes agent bug) + payload_copy = copy.deepcopy(payload) + # Deepcopy request_tags for _update_tag_db + request_tags = copy.deepcopy(payload.get("request_tags")) + + # Keep _insert_spend_log_to_db awaited inline (not a task, preserve current behavior) if disable_spend_logs is False: await self._insert_spend_log_to_db( payload=copy.deepcopy(payload), @@ -181,44 +148,20 @@ class DBSpendUpdateWriter: "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." ) + # Single task replaces 11 create_task() calls asyncio.create_task( - self.add_spend_log_transaction_to_daily_user_transaction( - payload=copy.deepcopy(payload), - prisma_client=prisma_client, - ) - ) - - asyncio.create_task( - self.add_spend_log_transaction_to_daily_end_user_transaction( - payload=copy.deepcopy(payload), - prisma_client=prisma_client, - ) - ) - - asyncio.create_task( - self.add_spend_log_transaction_to_daily_agent_transaction( - payload=payload, - prisma_client=prisma_client, - ) - ) - - asyncio.create_task( - self.add_spend_log_transaction_to_daily_team_transaction( - payload=copy.deepcopy(payload), - prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self.add_spend_log_transaction_to_daily_org_transaction( - payload=copy.deepcopy(payload), + self._batch_database_updates( + response_cost=response_cost, + user_id=user_id, + hashed_token=hashed_token, + team_id=team_id, org_id=org_id, + end_user_id=end_user_id, prisma_client=prisma_client, - ) - ) - asyncio.create_task( - self.add_spend_log_transaction_to_daily_tag_transaction( - payload=copy.deepcopy(payload), - prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + litellm_proxy_budget_name=litellm_proxy_budget_name, + payload_copy=payload_copy, + request_tags=request_tags, ) ) @@ -228,6 +171,157 @@ class DBSpendUpdateWriter: f"Error updating Prisma database: {traceback.format_exc()}" ) + async def _batch_database_updates( + self, + *, + response_cost: Optional[float], + user_id: Optional[str], + hashed_token: Optional[str], + team_id: Optional[str], + org_id: Optional[str], + end_user_id: Optional[str], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + litellm_proxy_budget_name: Optional[str], + payload_copy: dict, + request_tags: Optional[Any], + ): + """ + Runs all 11 spend-update helpers sequentially inside a single asyncio task. + + Each helper is wrapped in try/except so one failure doesn't prevent the others. + """ + try: + await self._update_user_db( + response_cost=response_cost, + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + litellm_proxy_budget_name=litellm_proxy_budget_name, + end_user_id=end_user_id, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_user_db failed: %s", + traceback.format_exc(), + ) + + try: + await self._update_key_db( + response_cost=response_cost, + hashed_token=hashed_token, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_key_db failed: %s", + traceback.format_exc(), + ) + + try: + await self._update_team_db( + response_cost=response_cost, + team_id=team_id, + user_id=user_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_team_db failed: %s", + traceback.format_exc(), + ) + + try: + await self._update_org_db( + response_cost=response_cost, + org_id=org_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_org_db failed: %s", + traceback.format_exc(), + ) + + try: + await self._update_tag_db( + response_cost=response_cost, + request_tags=request_tags, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_tag_db failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_user_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_user_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_end_user_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_agent_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_team_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_team_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_org_transaction( + payload=payload_copy, + org_id=org_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_org_transaction failed: %s", + traceback.format_exc(), + ) + + try: + await self.add_spend_log_transaction_to_daily_tag_transaction( + payload=payload_copy, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: add_spend_log_transaction_to_daily_tag_transaction failed: %s", + traceback.format_exc(), + ) + async def _update_key_db( self, response_cost: Optional[float], @@ -880,7 +974,7 @@ class DBSpendUpdateWriter: team_id = key.split("::")[1] user_id = key.split("::")[3] team_memberships_to_invalidate.append((user_id, team_id)) - + for i in range(n_retry_times + 1): start_time = time.time() try: @@ -917,11 +1011,13 @@ class DBSpendUpdateWriter: _raise_failed_update_spend_exception( e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) - + # Invalidate cache for updated team memberships # This ensures budget checks read fresh spend data from the database if team_memberships_to_invalidate and proxy_logging_obj is not None: - user_api_key_cache = proxy_logging_obj.call_details.get("user_api_key_cache") + user_api_key_cache = proxy_logging_obj.call_details.get( + "user_api_key_cache" + ) if user_api_key_cache is not None: for user_id, team_id in team_memberships_to_invalidate: cache_key = "team_membership:{}:{}".format(user_id, team_id) @@ -1233,7 +1329,9 @@ class DBSpendUpdateWriter: ), "endpoint": transaction.get("endpoint") or "", "prompt_tokens": transaction["prompt_tokens"], - "completion_tokens": transaction["completion_tokens"], + "completion_tokens": transaction[ + "completion_tokens" + ], "spend": transaction["spend"], "api_requests": transaction["api_requests"], "successful_requests": transaction[ @@ -1244,12 +1342,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: @@ -1292,10 +1392,14 @@ class DBSpendUpdateWriter: } if entity_type == "tag" and "request_id" in transaction: - update_data["request_id"] = transaction.get("request_id") + update_data["request_id"] = transaction.get( + "request_id" + ) # Add endpoint to update_data so existing rows get their endpoint field updated - update_data["endpoint"] = transaction.get("endpoint") or "" + update_data["endpoint"] = ( + transaction.get("endpoint") or "" + ) table.upsert( where=where_clause, @@ -1479,7 +1583,9 @@ class DBSpendUpdateWriter: self, payload: Union[dict, SpendLogsPayload], prisma_client: PrismaClient, - type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user", + type: Literal[ + "user", "team", "org", "request_tags", "end_user", "agent" + ] = "user", ) -> Optional[BaseDailySpendTransaction]: common_expected_keys = ["startTime", "api_key"] if type == "user": @@ -1538,7 +1644,7 @@ class DBSpendUpdateWriter: endpoint = None if call_type: endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None) - + daily_transaction = BaseDailySpendTransaction( date=date, api_key=payload["api_key"], @@ -1750,7 +1856,7 @@ class DBSpendUpdateWriter: endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}_{endpoint_str}" daily_transaction = DailyAgentSpendTransaction( - agent_id=payload['agent_id'], **base_daily_transaction + agent_id=payload["agent_id"], **base_daily_transaction ) await self.daily_agent_spend_update_queue.add_update( update={daily_transaction_key: daily_transaction} diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 0fa0d4cf10b..0465f2adaf3 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -51,6 +52,9 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): # Call the method await db_writer.update_database(**test_data) + # Let the single batched task run + await asyncio.sleep(0) + # Verify that _insert_spend_log_to_db was NOT called (since disable_spend_logs is True) db_writer._insert_spend_log_to_db.assert_not_called() @@ -115,7 +119,9 @@ async def test_update_daily_spend_with_null_entity_id(): # Verify the where clause contains null entity_id call_args = mock_table.upsert.call_args[1] - where_clause = call_args["where"]["user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"] + where_clause = call_args["where"][ + "user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" + ] assert where_clause["user_id"] is None assert where_clause["date"] == "2024-01-01" assert where_clause["api_key"] == "test-api-key" @@ -161,7 +167,7 @@ async def test_update_daily_spend_sorting(): upsert_calls = [] for i in range(50): daily_spend_transactions[f"test_key_{i}"] = { - "user_id": f"user{60-i}", # user60 ... user11, reverse order + "user_id": f"user{60-i}", # user60 ... user11, reverse order "date": "2024-01-01", "api_key": "test-api-key", "model": "gpt-4", @@ -173,46 +179,48 @@ async def test_update_daily_spend_sorting(): "successful_requests": 1, "failed_requests": 0, } - upsert_calls.append(call( - where={ - "user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": { - "user_id": f"user{i+11}", # user11 ... user60, sorted order - "date": "2024-01-01", - "api_key": "test-api-key", - "model": "gpt-4", - "custom_llm_provider": "openai", - "mcp_namespaced_tool_name": "", - "endpoint": "", - } - }, - data={ - "create": { - "user_id": f"user{i+11}", - "date": "2024-01-01", - "api_key": "test-api-key", - "model": "gpt-4", - "model_group": None, - "mcp_namespaced_tool_name": "", - "custom_llm_provider": "openai", - "endpoint": "", - "prompt_tokens": 10, - "completion_tokens": 20, - "spend": 0.1, - "api_requests": 1, - "successful_requests": 1, - "failed_requests": 0, + upsert_calls.append( + call( + where={ + "user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": { + "user_id": f"user{i+11}", # user11 ... user60, sorted order + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": "", + "endpoint": "", + } }, - "update": { - "prompt_tokens": {"increment": 10}, - "completion_tokens": {"increment": 20}, - "spend": {"increment": 0.1}, - "api_requests": {"increment": 1}, - "successful_requests": {"increment": 1}, - "failed_requests": {"increment": 0}, - "endpoint": "", + data={ + "create": { + "user_id": f"user{i+11}", + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "model_group": None, + "mcp_namespaced_tool_name": "", + "custom_llm_provider": "openai", + "endpoint": "", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + }, + "update": { + "prompt_tokens": {"increment": 10}, + "completion_tokens": {"increment": 20}, + "spend": {"increment": 0.1}, + "api_requests": {"increment": 1}, + "successful_requests": {"increment": 1}, + "failed_requests": {"increment": 0}, + "endpoint": "", + }, }, - }, - )) + ) + ) # Call the method await DBSpendUpdateWriter._update_daily_spend( @@ -275,7 +283,7 @@ async def test_update_daily_spend_tag_with_request_id(): # Verify that table.upsert was called mock_table.upsert.assert_called_once() - + # Verify request_id is in update_data call_args = mock_table.upsert.call_args[1] update_data = call_args["data"]["update"] @@ -283,15 +291,13 @@ async def test_update_daily_spend_tag_with_request_id(): assert update_data["request_id"] == "test-request-id-123" - - @pytest.mark.asyncio async def test_update_daily_spend_with_none_values_in_sorting_fields(): """ Test that _update_daily_spend handles None values in sorting fields correctly. - + This test ensures that when fields like date, api_key, model, or custom_llm_provider - are None, the sorting doesn't crash with TypeError: '<' not supported between + are None, the sorting doesn't crash with TypeError: '<' not supported between instances of 'NoneType' and 'str'. """ # Setup @@ -509,6 +515,7 @@ async def test_update_tag_db_without_prisma_client(): assert writer.spend_update_queue.add_update.call_count == 0 + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """ @@ -518,7 +525,7 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i writer = DBSpendUpdateWriter() mock_prisma = MagicMock() mock_prisma.get_request_status = MagicMock(return_value="success") - + request_id = "test-request-id-123" payload = { "request_id": request_id, @@ -546,13 +553,15 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i # Should be called twice (once for each tag) assert writer.daily_tag_spend_update_queue.add_update.call_count == 2 - + # Check that request_id is included in both transactions for call in writer.daily_tag_spend_update_queue.add_update.call_args_list: transaction_dict = call[1]["update"] # Each transaction should have one key with the format tag_date_api_key_model_provider for key, transaction in transaction_dict.items(): - assert transaction["request_id"] == request_id, f"request_id should be {request_id} but got {transaction.get('request_id')}" + assert ( + transaction["request_id"] == request_id + ), f"request_id should be {request_id} but got {transaction.get('request_id')}" @pytest.mark.asyncio @@ -866,11 +875,11 @@ async def test_endpoint_field_is_correctly_mapped_from_call_type(): call_args = writer.daily_spend_update_queue.add_update.call_args[1] update_dict = call_args["update"] assert len(update_dict) == 1 - + for key, transaction in update_dict.items(): # Verify endpoint is included in the key assert key == f"test-user_2024-01-01_test-key_gpt-4_openai_/chat/completions" - + # Verify endpoint is set in the transaction assert transaction["endpoint"] == "/chat/completions" assert transaction["user_id"] == "test-user" @@ -887,7 +896,7 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): This ensures proper debugging information is available for issues like unique constraint violations. """ from litellm._logging import verbose_proxy_logger - + # Setup mock_prisma_client = MagicMock() mock_batcher = MagicMock() @@ -895,13 +904,13 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): mock_batch_context = MagicMock() mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher) mock_batcher.litellm_dailyuserspend = mock_table - + # Make the batch context manager's exit raise an exception # This simulates a batch commit failure (e.g., unique constraint violation) test_exception = Exception("Unique constraint violation") mock_batch_context.__aexit__ = AsyncMock(side_effect=test_exception) mock_prisma_client.db.batch_.return_value = mock_batch_context - + # Create a transaction daily_spend_transactions = { "test_key": { @@ -918,13 +927,13 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): "failed_requests": 0, } } - + # Create a mock proxy_logging_obj with failure_handler as AsyncMock mock_proxy_logging = MagicMock() mock_proxy_logging.failure_handler = AsyncMock() - + # Mock the logger to capture exception calls - with patch.object(verbose_proxy_logger, 'exception') as mock_exception_logger: + with patch.object(verbose_proxy_logger, "exception") as mock_exception_logger: # Call the method and expect it to raise the exception with pytest.raises(Exception, match="Unique constraint violation"): await DBSpendUpdateWriter._update_daily_spend( @@ -937,13 +946,16 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): table_name="litellm_dailyuserspend", unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - + # Verify that exception was logged with detailed information assert mock_exception_logger.called call_args = mock_exception_logger.call_args[0][0] assert "Daily user spend batch upsert failed" in call_args assert "Table: litellm_dailyuserspend" in call_args - assert "Constraint: user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" in call_args + assert ( + "Constraint: user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" + in call_args + ) assert "Batch size: 1" in call_args assert "Unique constraint violation" in call_args @@ -961,7 +973,7 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): mock_batch_context = MagicMock() mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher) mock_batcher.litellm_dailyuserspend = mock_table - + # Create a transaction daily_spend_transactions = { "test_key": { @@ -978,16 +990,16 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): "failed_requests": 0, } } - + # Create a custom exception to verify it's re-raised custom_exception = ValueError("Database connection lost") mock_batch_context.__aexit__ = AsyncMock(side_effect=custom_exception) mock_prisma_client.db.batch_.return_value = mock_batch_context - + # Create a mock proxy_logging_obj with failure_handler as AsyncMock mock_proxy_logging = MagicMock() mock_proxy_logging.failure_handler = AsyncMock() - + # Verify the exception is re-raised with pytest.raises(ValueError, match="Database connection lost"): await DBSpendUpdateWriter._update_daily_spend( @@ -1018,10 +1030,12 @@ async def test_commit_key_spend_updates_includes_last_active(): mock_transaction = AsyncMock() mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) mock_transaction.__aexit__ = AsyncMock(return_value=False) - mock_transaction.batch_ = MagicMock(return_value=AsyncMock( - __aenter__=AsyncMock(return_value=mock_batcher), - __aexit__=AsyncMock(return_value=False), - )) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) mock_prisma_client = MagicMock() mock_prisma_client.db = MagicMock() @@ -1049,9 +1063,7 @@ async def test_commit_key_spend_updates_includes_last_active(): before_call = datetime.now(timezone.utc) - with patch( - "litellm.proxy.utils._raise_failed_update_spend_exception" - ): + with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): await db_writer._commit_spend_updates_to_db( prisma_client=mock_prisma_client, n_retry_times=0, @@ -1076,3 +1088,146 @@ async def test_commit_key_spend_updates_includes_last_active(): last_active = call_kwargs["data"]["last_active"] assert isinstance(last_active, datetime) assert before_call <= last_active <= after_call + + +@pytest.mark.asyncio +async def test_update_database_creates_single_task(): + """ + Test that update_database() fires exactly 1 asyncio.create_task() call + (the batched task) instead of the previous 11. + """ + db_writer = DBSpendUpdateWriter() + + # Mock all helpers so nothing real runs + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + + with patch("litellm.proxy.proxy_server.disable_spend_logs", False), patch( + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch( + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), patch( + "litellm.proxy.db.db_spend_update_writer.asyncio.create_task" + ) as mock_create_task: + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id="test-end-user", + start_time=datetime.now(), + end_time=datetime.now(), + team_id="test-team", + org_id="test-org", + completion_response=MagicMock(), + response_cost=0.1, + kwargs={"model": "gpt-4", "custom_llm_provider": "openai"}, + ) + + # Exactly 1 create_task call (the batch), not 11 + assert mock_create_task.call_count == 1 + + +@pytest.mark.asyncio +async def test_batch_database_updates_isolation_on_failure(): + """ + Test that if one helper inside _batch_database_updates raises, + all other helpers still execute. + """ + db_writer = DBSpendUpdateWriter() + + # Make _update_key_db raise + db_writer._update_key_db = AsyncMock(side_effect=RuntimeError("key db boom")) + + # All other helpers are normal mocks + db_writer._update_user_db = AsyncMock() + db_writer._update_team_db = AsyncMock() + db_writer._update_org_db = AsyncMock() + db_writer._update_tag_db = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() + + await db_writer._batch_database_updates( + response_cost=0.1, + user_id="u1", + hashed_token="t1", + team_id="team1", + org_id="org1", + end_user_id="eu1", + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + litellm_proxy_budget_name="budget", + payload_copy={"key": "value"}, + request_tags=None, + ) + + # _update_key_db raised, but all others should still have been called + db_writer._update_user_db.assert_awaited_once() + db_writer._update_key_db.assert_awaited_once() + db_writer._update_team_db.assert_awaited_once() + db_writer._update_org_db.assert_awaited_once() + db_writer._update_tag_db.assert_awaited_once() + db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once() + db_writer.add_spend_log_transaction_to_daily_end_user_transaction.assert_awaited_once() + db_writer.add_spend_log_transaction_to_daily_agent_transaction.assert_awaited_once() + db_writer.add_spend_log_transaction_to_daily_team_transaction.assert_awaited_once() + db_writer.add_spend_log_transaction_to_daily_org_transaction.assert_awaited_once() + db_writer.add_spend_log_transaction_to_daily_tag_transaction.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_daily_agent_receives_deepcopied_payload(): + """ + Test that the daily agent handler receives a deepcopied payload (not the original). + + Previously, add_spend_log_transaction_to_daily_agent_transaction received the raw + payload without a deepcopy, which was a mutation bug. + """ + db_writer = DBSpendUpdateWriter() + + original_payload = {"key": "value", "nested": {"a": 1}} + captured_payloads = [] + + async def capture_payload(**kwargs): + captured_payloads.append(kwargs.get("payload")) + + # Mock all helpers + db_writer._update_user_db = AsyncMock() + db_writer._update_key_db = AsyncMock() + db_writer._update_team_db = AsyncMock() + db_writer._update_org_db = AsyncMock() + db_writer._update_tag_db = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock( + side_effect=capture_payload + ) + db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() + + import copy + + payload_copy = copy.deepcopy(original_payload) + + await db_writer._batch_database_updates( + response_cost=0.1, + user_id="u1", + hashed_token="t1", + team_id="team1", + org_id="org1", + end_user_id="eu1", + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + litellm_proxy_budget_name="budget", + payload_copy=payload_copy, + request_tags=None, + ) + + # The payload passed to the agent handler should NOT be the original object + assert len(captured_payloads) == 1 + assert captured_payloads[0] is not original_payload + # But it should have the same content + assert captured_payloads[0] == original_payload From 1c48d8fda7512b386558044cec5bfbe18e22ee5c Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 24 Feb 2026 23:37:09 +0530 Subject: [PATCH 019/132] Add gpt-5.3-codex in model cost map --- ...odel_prices_and_context_window_backup.json | 33 +++++++++++++++++++ model_prices_and_context_window.json | 33 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3909ce4c8b0..00b960a9f7a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -20562,6 +20562,39 @@ "supports_tool_choice": true, "supports_vision": true }, + "gpt-5.3-codex": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": false, + "supports_tool_choice": true, + "supports_vision": true + }, "gpt-5-mini": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3909ce4c8b0..00b960a9f7a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -20562,6 +20562,39 @@ "supports_tool_choice": true, "supports_vision": true }, + "gpt-5.3-codex": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": false, + "supports_tool_choice": true, + "supports_vision": true + }, "gpt-5-mini": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, From c43a8dc84257f07ebbd2211a9f30e3fa8ddb92e5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 24 Feb 2026 10:17:35 -0800 Subject: [PATCH 020/132] feat(proxy): add warning/error level logging throughout spend tracking lifecycle Elevate silent debug-level and bare except:pass error paths to warning/error so spend tracking failures are visible in production logs. All new log messages are prefixed with "Spend tracking -" for easy filtering. Changes cover the full request-to-DB lifecycle: enqueue, in-memory flush, Redis buffer push/pop, DB commit, cache updates, spend log writes, and pod lock management. Also fixes a copy-paste bug in _update_team_cache that logged "end user" instead of "team". --- litellm/proxy/db/db_spend_update_writer.py | 95 ++++++++++++++++--- .../daily_spend_update_queue.py | 5 + .../db_transaction_queue/pod_lock_manager.py | 14 ++- .../redis_update_buffer.py | 38 ++++++-- .../spend_update_queue.py | 5 + litellm/proxy/proxy_server.py | 36 +++++-- litellm/proxy/utils.py | 19 +++- 7 files changed, 177 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 03628fda47f..c87be7ac819 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -224,8 +224,16 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("Runs spend update on all tables") except Exception: - verbose_proxy_logger.debug( - f"Error updating Prisma database: {traceback.format_exc()}" + verbose_proxy_logger.error( + "Spend tracking - update_database failed. All spend updates for this request will be lost. " + "response_cost=%s, token=%s, user_id=%s, team_id=%s, org_id=%s, end_user_id=%s - %s", + response_cost, + token, + user_id, + team_id, + org_id, + end_user_id, + traceback.format_exc(), ) async def _update_key_db( @@ -295,9 +303,14 @@ class DBSpendUpdateWriter: ) ) except Exception as e: - verbose_proxy_logger.debug( - "\033[91m" - + f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}" + verbose_proxy_logger.error( + "Spend tracking - failed to enqueue user spend update. " + "user_id=%s, end_user_id=%s, response_cost=%s - %s\n%s", + user_id, + end_user_id, + response_cost, + str(e), + traceback.format_exc(), ) async def _update_team_db( @@ -334,11 +347,23 @@ class DBSpendUpdateWriter: response_cost=response_cost, ) ) - except Exception: - pass + except Exception as e: + verbose_proxy_logger.error( + "Spend tracking - failed to enqueue team member spend update. " + "team_id=%s, user_id=%s, response_cost=%s - %s", + team_id, + user_id, + response_cost, + str(e), + ) except Exception as e: - verbose_proxy_logger.debug( - f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}" + verbose_proxy_logger.error( + "Spend tracking - failed to enqueue team spend update. " + "team_id=%s, response_cost=%s - %s\n%s", + team_id, + response_cost, + str(e), + traceback.format_exc(), ) raise e @@ -363,8 +388,13 @@ class DBSpendUpdateWriter: ) ) except Exception as e: - verbose_proxy_logger.debug( - f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}" + verbose_proxy_logger.error( + "Spend tracking - failed to enqueue org spend update. " + "org_id=%s, response_cost=%s - %s\n%s", + org_id, + response_cost, + str(e), + traceback.format_exc(), ) raise e @@ -411,8 +441,13 @@ class DBSpendUpdateWriter: ) ) except Exception as e: - verbose_proxy_logger.debug( - f"Update Tag DB failed to execute - {str(e)}\n{traceback.format_exc()}" + verbose_proxy_logger.error( + "Spend tracking - failed to enqueue tag spend update. " + "request_tags=%s, response_cost=%s - %s\n%s", + request_tags, + response_cost, + str(e), + traceback.format_exc(), ) raise e @@ -513,6 +548,17 @@ class DBSpendUpdateWriter: await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer() ) 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 {}), + ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, @@ -583,7 +629,12 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) except Exception as e: - verbose_proxy_logger.error(f"Error committing spend updates: {e}") + verbose_proxy_logger.error( + "Spend tracking - failed to commit spend updates from Redis to DB. " + "Data already popped from Redis may be lost. Error: %s\n%s", + str(e), + traceback.format_exc(), + ) finally: await self.pod_lock_manager.release_lock( cronjob_id=DB_SPEND_UPDATE_JOB_NAME, @@ -608,6 +659,22 @@ class DBSpendUpdateWriter: db_spend_update_transactions = ( await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() ) + if any( + len(v) > 0 + for v in db_spend_update_transactions.values() + if isinstance(v, dict) + ): + verbose_proxy_logger.info( + "Spend tracking - committing spend updates to DB (no Redis buffer): " + "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 {}), + ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index 5ba8fb13596..f47b694d44e 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -86,6 +86,11 @@ class DailySpendUpdateQueue(BaseUpdateQueue): ) -> Dict[str, BaseDailySpendTransaction]: """Get all updates from the queue and return all updates aggregated by daily_transaction_key. Works for both user and team spend updates.""" updates = await self.flush_all_updates_from_in_memory_queue() + if len(updates) > 0: + verbose_proxy_logger.info( + "Spend tracking - flushed %d daily spend update items from in-memory queue", + len(updates), + ) aggregated_daily_spend_update_transactions = ( DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( updates diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index bb5424b0e90..6f86e82cf29 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -80,6 +80,14 @@ class PodLockManager: ) self._emit_acquired_lock_event(cronjob_id, self.pod_id) return True + else: + verbose_proxy_logger.info( + "Spend tracking - pod %s could not acquire lock for cronjob_id=%s, " + "held by pod %s. Spend updates in Redis will wait for the leader pod to commit.", + self.pod_id, + cronjob_id, + current_value, + ) return False except Exception as e: verbose_proxy_logger.error( @@ -124,10 +132,12 @@ class PodLockManager: pod_id=self.pod_id, ) else: - verbose_proxy_logger.debug( - "Pod %s failed to release Redis lock for cronjob_id=%s", + verbose_proxy_logger.warning( + "Spend tracking - pod %s failed to release Redis lock for cronjob_id=%s. " + "Lock will expire after TTL=%ds.", self.pod_id, cronjob_id, + DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, ) else: verbose_proxy_logger.debug( diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 37b42e26bc9..85e139b4de0 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -96,14 +96,29 @@ class RedisUpdateBuffer: list_of_transactions = [safe_dumps(transactions)] if self.redis_cache is None: return - current_redis_buffer_size = await self.redis_cache.async_rpush( - key=redis_key, - values=list_of_transactions, - ) - await self._emit_new_item_added_to_redis_buffer_event( - queue_size=current_redis_buffer_size, - service=service_type, - ) + try: + current_redis_buffer_size = await self.redis_cache.async_rpush( + key=redis_key, + values=list_of_transactions, + ) + verbose_proxy_logger.info( + "Spend tracking - pushed spend updates to Redis buffer. " + "redis_key=%s, buffer_size=%s", + redis_key, + current_redis_buffer_size, + ) + await self._emit_new_item_added_to_redis_buffer_event( + queue_size=current_redis_buffer_size, + service=service_type, + ) + except Exception as e: + verbose_proxy_logger.error( + "Spend tracking - failed to push spend updates to Redis (redis_key=%s). " + "Error: %s", + redis_key, + str(e), + ) + raise async def store_in_memory_spend_updates_in_redis( self, @@ -305,6 +320,13 @@ class RedisUpdateBuffer: if list_of_transactions is None: return None + verbose_proxy_logger.info( + "Spend tracking - popped %d spend update batches from Redis buffer (key=%s). " + "These items are now removed from Redis and must be committed to DB.", + len(list_of_transactions) if isinstance(list_of_transactions, list) else 1, + REDIS_UPDATE_BUFFER_KEY, + ) + # Parse the list of transactions from JSON strings parsed_transactions = self._parse_list_of_transactions(list_of_transactions) diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index b41ff121622..3e059cf8c1f 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -31,6 +31,11 @@ class SpendUpdateQueue(BaseUpdateQueue): ) -> DBSpendUpdateTransactions: """Flush all updates from the queue and return all updates aggregated by entity type.""" updates = await self.flush_all_updates_from_in_memory_queue() + if len(updates) > 0: + verbose_proxy_logger.info( + "Spend tracking - flushed %d spend update items from in-memory queue", + len(updates), + ) verbose_proxy_logger.debug("Aggregating updates by entity type: %s", updates) return self.get_aggregated_db_spend_update_transactions(updates) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b3b6bf0ccf7..78730ee9d60 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1768,8 +1768,13 @@ async def update_cache( # noqa: PLR0915 ("{}:spend".format(litellm_proxy_admin_name), increment) ) except Exception as e: - verbose_proxy_logger.debug( - f"An error occurred updating user cache: {str(e)}\n\n{traceback.format_exc()}" + verbose_proxy_logger.warning( + "Spend tracking - failed to update user spend in cache. " + "Budget enforcement may use stale spend values. " + "user_id=%s, response_cost=%s - %s", + user_id, + response_cost, + str(e), ) ### UPDATE END-USER SPEND ### @@ -1806,8 +1811,13 @@ async def update_cache( # noqa: PLR0915 existing_spend_obj.spend = new_spend values_to_update_in_cache.append((_id, existing_spend_obj.json())) except Exception as e: - verbose_proxy_logger.exception( - f"An error occurred updating end user cache: {str(e)}" + verbose_proxy_logger.warning( + "Spend tracking - failed to update end user spend in cache. " + "Budget enforcement may use stale spend values. " + "end_user_id=%s, response_cost=%s - %s", + end_user_id, + response_cost, + str(e), ) ### UPDATE TEAM SPEND ### @@ -1848,8 +1858,13 @@ async def update_cache( # noqa: PLR0915 existing_spend_obj.spend = new_spend values_to_update_in_cache.append((_id, existing_spend_obj)) except Exception as e: - verbose_proxy_logger.exception( - f"An error occurred updating end user cache: {str(e)}" + verbose_proxy_logger.warning( + "Spend tracking - failed to update team spend in cache. " + "Budget enforcement may use stale spend values. " + "team_id=%s, response_cost=%s - %s", + team_id, + response_cost, + str(e), ) ### UPDATE TAG SPEND ### @@ -1894,8 +1909,13 @@ async def update_cache( # noqa: PLR0915 existing_tag_obj.spend = new_spend values_to_update_in_cache.append((cache_key, existing_tag_obj)) except Exception as e: - verbose_proxy_logger.exception( - f"An error occurred updating tag cache: {str(e)}" + verbose_proxy_logger.warning( + "Spend tracking - failed to update tag spend in cache. " + "Budget enforcement may use stale spend values. " + "tags=%s, response_cost=%s - %s", + tags, + response_cost, + str(e), ) if token is not None and response_cost is not None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c4ff325db1f..f6613b5548f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4457,6 +4457,11 @@ class ProxyUpdateSpend: len(logs_to_process) : ] popped_batch = True + if len(logs_to_process) > 0: + verbose_proxy_logger.info( + "Spend tracking - processing %d spend logs for DB write", + len(logs_to_process), + ) start_time = time.time() try: for i in range(n_retry_times + 1): @@ -4503,9 +4508,17 @@ class ProxyUpdateSpend: f"{len(logs_to_process)} logs processed. Remaining in queue: {remaining_count}" ) break - except DB_CONNECTION_ERROR_TYPES: + except DB_CONNECTION_ERROR_TYPES as e: if i is None: i = 0 + verbose_proxy_logger.warning( + "Spend tracking - DB connection error writing spend logs, " + "retry %d/%d. logs_count=%d, error=%s", + i + 1, + n_retry_times, + len(logs_to_process), + str(e), + ) if i >= n_retry_times: raise await asyncio.sleep(2**i) @@ -4620,8 +4633,8 @@ async def update_spend_logs_job( logs_to_process=logs_to_process, ) except Exception as guardrail_tracking_err: - verbose_proxy_logger.debug( - "Guardrail usage tracking failed (non-fatal): %s", + verbose_proxy_logger.warning( + "Spend tracking - guardrail usage tracking failed (non-fatal): %s", guardrail_tracking_err, ) From 4c5963fdb90fd33b64cb4687a8ce8c41686810e7 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 24 Feb 2026 10:49:16 -0800 Subject: [PATCH 021/132] test: exercise production deepcopy path in agent payload test --- .../proxy/db/test_db_spend_update_writer.py | 75 ++++++++++++------- 1 file changed, 49 insertions(+), 26 deletions(-) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 0465f2adaf3..add7293285a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1183,17 +1183,21 @@ async def test_daily_agent_receives_deepcopied_payload(): Test that the daily agent handler receives a deepcopied payload (not the original). Previously, add_spend_log_transaction_to_daily_agent_transaction received the raw - payload without a deepcopy, which was a mutation bug. + payload without a deepcopy, which was a mutation bug. This test goes through + update_database() to verify the production deepcopy path. """ db_writer = DBSpendUpdateWriter() - original_payload = {"key": "value", "nested": {"a": 1}} - captured_payloads = [] + # Capture the payload object that get_logging_payload returns (the "original") + # and the payload the agent handler receives (should be a deepcopy) + original_payload_ref = {} + captured_agent_payloads = [] - async def capture_payload(**kwargs): - captured_payloads.append(kwargs.get("payload")) + async def capture_agent_payload(**kwargs): + captured_agent_payloads.append(kwargs.get("payload")) # Mock all helpers + db_writer._insert_spend_log_to_db = AsyncMock() db_writer._update_user_db = AsyncMock() db_writer._update_key_db = AsyncMock() db_writer._update_team_db = AsyncMock() @@ -1202,32 +1206,51 @@ async def test_daily_agent_receives_deepcopied_payload(): db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock( - side_effect=capture_payload + side_effect=capture_agent_payload ) db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() - import copy + # Mock get_logging_payload to return a known dict and capture its identity + fake_payload = { + "startTime": "2024-01-01T00:00:00", + "endTime": "2024-01-01T00:01:00", + "model": "gpt-4", + "custom_llm_provider": "openai", + "spend": 0.0, + "nested": {"a": 1}, + } + original_payload_ref["obj"] = fake_payload # store reference to the original - payload_copy = copy.deepcopy(original_payload) + with patch("litellm.proxy.proxy_server.disable_spend_logs", True), patch( + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch( + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=fake_payload, + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id="test-end-user", + team_id="test-team", + org_id="test-org", + kwargs={"model": "gpt-4", "custom_llm_provider": "openai"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) - await db_writer._batch_database_updates( - response_cost=0.1, - user_id="u1", - hashed_token="t1", - team_id="team1", - org_id="org1", - end_user_id="eu1", - prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), - litellm_proxy_budget_name="budget", - payload_copy=payload_copy, - request_tags=None, - ) + # Let the single batched task run + await asyncio.sleep(0) - # The payload passed to the agent handler should NOT be the original object - assert len(captured_payloads) == 1 - assert captured_payloads[0] is not original_payload - # But it should have the same content - assert captured_payloads[0] == original_payload + # The agent handler should have been called + assert len(captured_agent_payloads) == 1 + # The payload must NOT be the same object as the original (deepcopy occurred) + assert captured_agent_payloads[0] is not original_payload_ref["obj"] + # But it should have equivalent content + assert captured_agent_payloads[0]["model"] == "gpt-4" + assert captured_agent_payloads[0]["spend"] == 0.1 From aded14a55ac973c91559e49fcef4d267cbc17229 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 25 Feb 2026 01:04:12 +0530 Subject: [PATCH 022/132] Fix release version for gpt-5.3-codex --- docs/my-website/blog/gpt_5_3_codex/index.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/my-website/blog/gpt_5_3_codex/index.md b/docs/my-website/blog/gpt_5_3_codex/index.md index 321e573cccf..5c43d287bce 100644 --- a/docs/my-website/blog/gpt_5_3_codex/index.md +++ b/docs/my-website/blog/gpt_5_3_codex/index.md @@ -44,7 +44,7 @@ Important: ## Docker Image ```bash -docker pull ghcr.io/berriai/litellm:v1.81.3-stable.sonnet-4-6 +docker pull ghcr.io/berriai/litellm:v1.81.12-stable.gpt-5.3 ``` ## Usage @@ -68,7 +68,7 @@ docker run -d \ -p 4000:4000 \ -e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \ -v $(pwd)/config.yaml:/app/config.yaml \ - ghcr.io/berriai/litellm:v1.81.3-stable.sonnet-4-6 \ + ghcr.io/berriai/litellm:v1.81.12-stable.gpt-5.3 \ --config /app/config.yaml ``` From 74abf0c8e6b6144032dc386972626b2b6ee22a98 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 25 Feb 2026 01:19:10 +0530 Subject: [PATCH 023/132] Fix phase docs link --- docs/my-website/blog/gpt_5_3_codex/index.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/my-website/blog/gpt_5_3_codex/index.md b/docs/my-website/blog/gpt_5_3_codex/index.md index 5c43d287bce..5b7773cf8aa 100644 --- a/docs/my-website/blog/gpt_5_3_codex/index.md +++ b/docs/my-website/blog/gpt_5_3_codex/index.md @@ -29,7 +29,7 @@ LiteLLM now supports GPT-5.3-Codex on Day 0, including support for the new assis `phase` appears on assistant output items and helps distinguish preamble/commentary turns from final closeout responses. -Reference: [Phase parameter docs](https://developers-site-git-alphas-venusaur-api-openai.vercel.app/alphas/venusaur-api/phase-parameter) +Reference: [Phase parameter docs](https://developers.openai.com/api/reference/overview) Supported values: - `null` From 5d291c739fbc493e3d2dae593add5ba6eb221c0a Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 25 Feb 2026 01:21:38 +0530 Subject: [PATCH 024/132] Fix phase docs link --- docs/my-website/blog/gpt_5_3_codex/index.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/my-website/blog/gpt_5_3_codex/index.md b/docs/my-website/blog/gpt_5_3_codex/index.md index 5b7773cf8aa..850586538f6 100644 --- a/docs/my-website/blog/gpt_5_3_codex/index.md +++ b/docs/my-website/blog/gpt_5_3_codex/index.md @@ -66,7 +66,7 @@ model_list: ```bash docker run -d \ -p 4000:4000 \ - -e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \ + -e ANTHROPIC_API_KEY=$OPENAI_API_KEY \ -v $(pwd)/config.yaml:/app/config.yaml \ ghcr.io/berriai/litellm:v1.81.12-stable.gpt-5.3 \ --config /app/config.yaml From e44b9b6b3584710a948365268d64676e2177adab Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 11:51:42 -0800 Subject: [PATCH 025/132] feat(prometheus): add opt-in stream label to litellm_proxy_total_requests_metric (#22023) Set prometheus_emit_stream_label: true in litellm_settings to emit a stream label (True/False/None) on litellm_proxy_total_requests_metric. Opt-in to avoid breaking cardinality on existing deployments. --- docs/my-website/docs/proxy/prometheus.md | 26 +++++- litellm/__init__.py | 1 + litellm/integrations/prometheus.py | 6 ++ litellm/types/integrations/prometheus.py | 25 ++++-- .../test_prometheus_stream_label.py | 81 +++++++++++++++++++ 5 files changed, 130 insertions(+), 9 deletions(-) create mode 100644 tests/test_litellm/integrations/test_prometheus_stream_label.py diff --git a/docs/my-website/docs/proxy/prometheus.md b/docs/my-website/docs/proxy/prometheus.md index 93a0675f097..18a139d1d29 100644 --- a/docs/my-website/docs/proxy/prometheus.md +++ b/docs/my-website/docs/proxy/prometheus.md @@ -122,7 +122,7 @@ Use this to track overall LiteLLM Proxy usage. | Metric Name | Description | |----------------------|--------------------------------------| | `litellm_proxy_failed_requests_metric` | Total number of failed responses from proxy - the client did not get a success response from litellm proxy. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "user_email", "exception_status", "exception_class", "route", "model_id"` | -| `litellm_proxy_total_requests_metric` | Total number of requests made to the proxy server - track number of client side requests. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "status_code", "user_email", "route", "model_id"` | +| `litellm_proxy_total_requests_metric` | Total number of requests made to the proxy server - track number of client side requests. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "status_code", "user_email", "route", "model_id"`. Optionally includes `"stream"` — see [Emit Stream Label](#emit-stream-label). | ### Callback Logging Metrics @@ -214,9 +214,31 @@ litellm_settings: ``` +### Emit Stream Label + +Add a `stream` label to `litellm_proxy_total_requests_metric` to split requests by streaming vs. non-streaming. Disabled by default. + +```yaml title="config.yaml" +litellm_settings: + callbacks: ["prometheus"] + prometheus_emit_stream_label: true +``` + +When enabled, `litellm_proxy_total_requests_metric` gains a `stream` label with values `"True"`, `"False"`, or `"None"`. + +``` +litellm_proxy_total_requests_metric{..., stream="True"} 42 +litellm_proxy_total_requests_metric{..., stream="False"} 100 +``` + +:::note +This label is opt-in because adding a new label to an existing metric changes its cardinality and breaks existing Prometheus queries / Grafana dashboards that target this metric. Enable it only on fresh deployments or when you are ready to update your dashboards. +::: + + ## [BETA] Custom Metrics -Track custom metrics on prometheus on all events mentioned above. +Track custom metrics on prometheus on all events mentioned above. ### Custom Metadata Labels diff --git a/litellm/__init__.py b/litellm/__init__.py index 1e74b5692e4..6e42f2c1ea5 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -374,6 +374,7 @@ enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None custom_prometheus_metadata_labels: List[str] = [] custom_prometheus_tags: List[str] = [] prometheus_metrics_config: Optional[List] = None +prometheus_emit_stream_label: bool = False disable_add_prefix_to_prompt: bool = ( False # used by anthropic, to disable adding prefix to prompt ) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 4c7afd5a57c..08db77e8571 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -974,6 +974,9 @@ class PrometheusLogger(CustomLogger): ), client_ip=standard_logging_payload["metadata"].get("requester_ip_address"), user_agent=standard_logging_payload["metadata"].get("user_agent"), + stream=str(standard_logging_payload.get("stream")) + if litellm.prometheus_emit_stream_label + else None, ) if ( @@ -1624,6 +1627,9 @@ class PrometheusLogger(CustomLogger): client_ip=_metadata.get("requester_ip_address"), user_agent=_metadata.get("user_agent"), model_id=model_id, + stream=str(request_data.get("stream")) + if litellm.prometheus_emit_stream_label + else None, ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index fd788af9ac1..482b87085dd 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -55,22 +55,21 @@ def _sanitize_prometheus_label_value(value: Optional[Any]) -> Optional[str]: return None # Coerce non-string values (int, bool, etc.) to str before sanitizing - if not isinstance(value, str): - value = str(value) + str_value: str = value if isinstance(value, str) else str(value) # Remove Unicode line/paragraph separators that break text format - value = value.replace("\u2028", "").replace("\u2029", "") + str_value = str_value.replace("\u2028", "").replace("\u2029", "") # Remove carriage returns - value = value.replace("\r", "") + str_value = str_value.replace("\r", "") # Replace newlines with spaces - value = value.replace("\n", " ") + str_value = str_value.replace("\n", " ") # Escape backslashes and double quotes per Prometheus exposition format - value = value.replace("\\", "\\\\").replace('"', '\\"') + str_value = str_value.replace("\\", "\\\\").replace('"', '\\"') - return value + return str_value @dataclass @@ -185,6 +184,7 @@ class UserAPIKeyLabelNames(Enum): CLIENT_IP = "client_ip" USER_AGENT = "user_agent" CALLBACK_NAME = "callback_name" + STREAM = "stream" DEFINED_PROMETHEUS_METRICS = Literal[ @@ -638,6 +638,14 @@ class PrometheusMetricLabels: ] ) + # Conditionally add stream label to litellm_proxy_total_requests_metric + if ( + label_name == "litellm_proxy_total_requests_metric" + and litellm.prometheus_emit_stream_label is True + and UserAPIKeyLabelNames.STREAM.value not in default_labels + ): + custom_labels.append(UserAPIKeyLabelNames.STREAM.value) + return default_labels + custom_labels @@ -709,6 +717,9 @@ class UserAPIKeyLabelValues(BaseModel): user_agent: Annotated[ Optional[str], Field(..., alias=UserAPIKeyLabelNames.USER_AGENT.value) ] = None + stream: Annotated[ + Optional[str], Field(..., alias=UserAPIKeyLabelNames.STREAM.value) + ] = None class PrometheusMetricsConfig(BaseModel): diff --git a/tests/test_litellm/integrations/test_prometheus_stream_label.py b/tests/test_litellm/integrations/test_prometheus_stream_label.py new file mode 100644 index 00000000000..a00a468e0fb --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_stream_label.py @@ -0,0 +1,81 @@ +""" +Unit tests for prometheus_emit_stream_label opt-in setting. + +Tests that: +- stream label is NOT added to litellm_proxy_total_requests_metric by default +- stream label IS added when litellm.prometheus_emit_stream_label = True +- stream value is populated correctly from standard_logging_payload +""" +import pytest + +import litellm +from litellm.types.integrations.prometheus import ( + PrometheusMetricLabels, + UserAPIKeyLabelNames, +) + + +def test_stream_label_not_present_by_default(): + """stream label should NOT appear in litellm_proxy_total_requests_metric unless opted in""" + litellm.prometheus_emit_stream_label = False + labels = PrometheusMetricLabels.get_labels("litellm_proxy_total_requests_metric") + assert UserAPIKeyLabelNames.STREAM.value not in labels + + +def test_stream_label_present_when_opted_in(): + """stream label SHOULD appear in litellm_proxy_total_requests_metric when opted in""" + litellm.prometheus_emit_stream_label = True + try: + labels = PrometheusMetricLabels.get_labels("litellm_proxy_total_requests_metric") + assert UserAPIKeyLabelNames.STREAM.value in labels + finally: + litellm.prometheus_emit_stream_label = False + + +def test_stream_label_not_in_other_metrics_when_opted_in(): + """stream label should NOT be added to other metrics even when opted in""" + litellm.prometheus_emit_stream_label = True + try: + other_metrics = [ + "litellm_proxy_failed_requests_metric", + "litellm_spend_metric", + "litellm_input_tokens_metric", + "litellm_output_tokens_metric", + "litellm_llm_api_latency_metric", + ] + for metric in other_metrics: + labels = PrometheusMetricLabels.get_labels(metric) + assert UserAPIKeyLabelNames.STREAM.value not in labels, ( + f"stream label should not be in {metric}" + ) + finally: + litellm.prometheus_emit_stream_label = False + + +def test_stream_label_name(): + """STREAM label name should be 'stream'""" + assert UserAPIKeyLabelNames.STREAM.value == "stream" + + +def test_user_api_key_label_values_has_stream_field(): + """UserAPIKeyLabelValues should accept stream field""" + from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + values = UserAPIKeyLabelValues(stream="True") + assert values.stream == "True" + + values_false = UserAPIKeyLabelValues(stream="False") + assert values_false.stream == "False" + + values_none = UserAPIKeyLabelValues() + assert values_none.stream is None + + +def test_stream_label_in_model_dump(): + """stream field appears in model_dump() output for use in prometheus_label_factory""" + from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + values = UserAPIKeyLabelValues(stream="True") + dumped = values.model_dump() + assert "stream" in dumped + assert dumped["stream"] == "True" From c343bfffdaa3e34e9324733d897a8e4db79ac9ec Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 11:56:16 -0800 Subject: [PATCH 026/132] fix(router): emit x-litellm-overhead-duration-ms header for streaming requests (#22027) * fix(router): preserve _hidden_params in FallbackStreamWrapper so x-litellm-overhead-duration-ms is emitted for streaming requests * test(router): add regression test for FallbackStreamWrapper _hidden_params preservation --- litellm/router.py | 26 ++- .../proxy/test_common_request_processing.py | 201 ++++++++++++++++++ tests/test_litellm/test_router.py | 57 ++++- 3 files changed, 274 insertions(+), 10 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index ac2862da689..46d35352c37 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -163,7 +163,11 @@ from litellm.types.utils import ( ) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage +from litellm.types.utils import ( + ModelResponseStream, + StandardLoggingPayload, + Usage, +) from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -1555,6 +1559,9 @@ class Router: logging_obj=model_response.logging_obj, ) self._async_generator = async_generator + # Preserve hidden params (including litellm_overhead_time_ms) from original response + if hasattr(model_response, "_hidden_params"): + self._hidden_params = model_response._hidden_params.copy() def __aiter__(self): return self @@ -6978,9 +6985,9 @@ class Router: raise ValueError("Deployment not found") ## GET BASE MODEL - base_model = deployment.get("model_info", {}).get("base_model", None) + base_model = (deployment.get("model_info") or {}).get("base_model", None) if base_model is None: - base_model = deployment.get("litellm_params", {}).get("base_model", None) + base_model = (deployment.get("litellm_params") or {}).get("base_model", None) model = base_model @@ -6995,7 +7002,7 @@ class Router: raise ValueError( f"Deployment missing valid litellm_params. " f"Got: {type(litellm_params_data).__name__}, " - f"deployment_id: {deployment.get('model_info', {}).get('id', 'unknown')}" + f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}" ) _model, custom_llm_provider, _, _ = litellm.get_llm_provider( model=litellm_params.model, @@ -7015,10 +7022,10 @@ class Router: if potential_models is not None: for potential_model in potential_models: try: - if potential_model.get("model_info", {}).get( + if (potential_model.get("model_info") or {}).get( "id" - ) == deployment.get("model_info", {}).get("id"): - model = potential_model.get("litellm_params", {}).get( + ) == (deployment.get("model_info") or {}).get("id"): + model = (potential_model.get("litellm_params") or {}).get( "model" ) break @@ -7039,9 +7046,10 @@ class Router: model_info = litellm.get_model_info(model=model_info_name) ## CHECK USER SET MODEL INFO - user_model_info = deployment.get("model_info", {}) + user_model_info = deployment.get("model_info") or {} - model_info.update(user_model_info) + if model_info is not None: + model_info.update(user_model_info) return model_info diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 977304f732b..bf794478f10 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,4 +1,6 @@ import copy +import datetime +from typing import AsyncGenerator from unittest.mock import AsyncMock, MagicMock import pytest @@ -1348,3 +1350,202 @@ class TestOverrideOpenAIResponseModel: # Verify the model was not changed assert response_obj.model == fallback_model + + +class TestStreamingOverheadHeader: + """ + Tests that x-litellm-overhead-duration-ms is emitted in streaming responses. + + Regression tests for: streaming requests not including overhead header. + """ + + def test_get_custom_headers_includes_overhead_when_set(self): + """ + get_custom_headers() returns x-litellm-overhead-duration-ms + when litellm_overhead_time_ms is in hidden_params. + """ + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + hidden_params = { + "litellm_overhead_time_ms": 42.5, + "_response_ms": 500.0, + "model_id": "test-model-id", + "api_base": "https://api.openai.com", + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + model_id="test-model-id", + cache_key="", + api_base="https://api.openai.com", + version="1.0.0", + response_cost=0.001, + model_region="", + hidden_params=hidden_params, + ) + + assert "x-litellm-overhead-duration-ms" in headers + assert headers["x-litellm-overhead-duration-ms"] == "42.5" + + def test_get_custom_headers_omits_overhead_when_none(self): + """ + get_custom_headers() omits x-litellm-overhead-duration-ms + when litellm_overhead_time_ms is not in hidden_params. + """ + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + hidden_params = { + "_response_ms": 500.0, + "model_id": "test-model-id", + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + model_id="test-model-id", + cache_key="", + api_base="https://api.openai.com", + version="1.0.0", + response_cost=0.001, + model_region="", + hidden_params=hidden_params, + ) + + # Should be absent (None gets filtered by exclude_values) + assert "x-litellm-overhead-duration-ms" not in headers + + def test_update_response_metadata_sets_overhead_on_stream_wrapper(self): + """ + update_response_metadata() sets litellm_overhead_time_ms on + a streaming response's _hidden_params when llm_api_duration_ms is available. + """ + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata, + ) + + # Mock the logging object with llm_api_duration_ms set + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = { + "llm_api_duration_ms": 200.0, + "litellm_params": {}, + } + mock_logging_obj.caching_details = None + mock_logging_obj.callback_duration_ms = None + mock_logging_obj.litellm_call_id = "test-call-id" + mock_logging_obj._response_cost_calculator = MagicMock(return_value=0.001) + + # Simulate a streaming result object with _hidden_params (like CustomStreamWrapper) + stream_result = MagicMock() + stream_result._hidden_params = { + "model_id": "test-model-id", + "api_base": "https://api.openai.com", + "additional_headers": {}, + } + + start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=300) + end_time = datetime.datetime.now() + + update_response_metadata( + result=stream_result, + logging_obj=mock_logging_obj, + model="gpt-4o", + kwargs={}, + start_time=start_time, + end_time=end_time, + ) + + assert "litellm_overhead_time_ms" in stream_result._hidden_params + overhead = stream_result._hidden_params["litellm_overhead_time_ms"] + assert overhead is not None + assert isinstance(overhead, float) + # overhead = total_response_ms (~300ms) - llm_api_duration_ms (200ms) = ~100ms + assert overhead > 0 + + @pytest.mark.asyncio + async def test_streaming_response_includes_overhead_header(self): + """ + StreamingResponse returned by create_response() includes + x-litellm-overhead-duration-ms in its headers. + """ + + async def mock_generator() -> AsyncGenerator[str, None]: + yield 'data: {"id":"chatcmpl-test","choices":[{"delta":{"content":"hi"}}]}\n\n' + yield "data: [DONE]\n\n" + + headers = { + "x-litellm-overhead-duration-ms": "42.5", + "x-litellm-call-id": "test-call-id", + "x-litellm-model-id": "test-model-id", + } + + response = await create_response( + generator=mock_generator(), + media_type="text/event-stream", + headers=headers, + ) + + assert isinstance(response, StreamingResponse) + assert response.headers.get("x-litellm-overhead-duration-ms") == "42.5" + + def test_streaming_overhead_header_in_custom_headers_from_stream_hidden_params( + self, + ): + """ + Verifies that when get_custom_headers() is called with a streaming + response's hidden_params (containing litellm_overhead_time_ms), + the x-litellm-overhead-duration-ms header is correctly populated. + + This tests the critical path: update_response_metadata sets the value + → get_custom_headers reads it → StreamingResponse header is set. + """ + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + # This is what CustomStreamWrapper._hidden_params looks like after + # update_response_metadata() has been called on it + hidden_params = { + "model_id": "openai-gpt4o-deployment", + "api_base": "https://api.openai.com", + "additional_headers": {}, + "litellm_overhead_time_ms": 55.3, # set by update_response_metadata + "_response_ms": 280.0, + "litellm_call_id": "test-call-id", + "response_cost": 0.002, + "cache_key": None, + "fastest_response_batch_completion": None, + "callback_duration_ms": None, + } + + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + model_id=hidden_params.get("model_id"), + cache_key=hidden_params.get("cache_key") or "", + api_base=hidden_params.get("api_base") or "", + version="1.0.0", + response_cost=hidden_params.get("response_cost"), + model_region="", + hidden_params=hidden_params, + ) + + # The overhead header must be present and correct + assert "x-litellm-overhead-duration-ms" in custom_headers, ( + "x-litellm-overhead-duration-ms header must be emitted during streaming. " + "It was missing — this is the streaming overhead header regression." + ) + assert custom_headers["x-litellm-overhead-duration-ms"] == "55.3" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5542cdf8be4..5732deda6fb 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1297,6 +1297,61 @@ async def test_acompletion_streaming_iterator_edge_cases(): print("✓ Edge case tests passed!") +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_preserves_hidden_params(): + """ + Regression test: FallbackStreamWrapper must copy _hidden_params from the + original CustomStreamWrapper so that x-litellm-overhead-duration-ms (and + other hidden params) are present in the proxy response headers for streaming. + """ + from unittest.mock import MagicMock + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + # Simulate a CustomStreamWrapper that already has timing metadata set by + # update_response_metadata (litellm_overhead_time_ms, _response_ms, etc.) + mock_response = MagicMock() + mock_response.model = "gpt-4" + mock_response.custom_llm_provider = "openai" + mock_response.logging_obj = MagicMock() + mock_response._hidden_params = { + "litellm_overhead_time_ms": 12.34, + "_response_ms": 500.0, + "litellm_call_id": "test-call-id", + "api_base": "https://api.openai.com", + "additional_headers": {}, + } + + # Make the mock iterable (yields nothing — we only care about hidden_params) + async def _empty(): + return + yield # make it an async generator + + mock_response.__aiter__ = lambda self: _empty().__aiter__() + + result = await router._acompletion_streaming_iterator( + model_response=mock_response, + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + + # The returned FallbackStreamWrapper must carry the original _hidden_params + assert hasattr(result, "_hidden_params"), "result must have _hidden_params" + assert result._hidden_params.get("litellm_overhead_time_ms") == 12.34, ( + "litellm_overhead_time_ms must be preserved — " + "this is what drives x-litellm-overhead-duration-ms in streaming responses" + ) + assert result._hidden_params.get("litellm_call_id") == "test-call-id" + assert result._hidden_params.get("_response_ms") == 500.0 + + @pytest.mark.asyncio async def test_async_function_with_fallbacks_common_utils(): """Test the async_function_with_fallbacks_common_utils method""" @@ -1858,7 +1913,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): litellm_credential_name to actual credential values (for UI-created models). """ from litellm.types.utils import CredentialItem - + # Setup credential list with a test credential litellm.credential_list = [ CredentialItem( From 235d60eb885210aca0d981a63ccfe708caa8ed81 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 24 Feb 2026 11:58:59 -0800 Subject: [PATCH 027/132] address greptile review feedback (greploop iteration 1) - Add traceback to cache update warning logs (user, end_user, team, tag) - Remove duplicate info log in non-redis commit path --- litellm/proxy/db/db_spend_update_writer.py | 16 ---------------- litellm/proxy/proxy_server.py | 12 ++++++++---- 2 files changed, 8 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c87be7ac819..e2d365ee465 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -659,22 +659,6 @@ class DBSpendUpdateWriter: db_spend_update_transactions = ( await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() ) - if any( - len(v) > 0 - for v in db_spend_update_transactions.values() - if isinstance(v, dict) - ): - verbose_proxy_logger.info( - "Spend tracking - committing spend updates to DB (no Redis buffer): " - "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 {}), - ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 78730ee9d60..1983a601e13 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1771,10 +1771,11 @@ async def update_cache( # noqa: PLR0915 verbose_proxy_logger.warning( "Spend tracking - failed to update user spend in cache. " "Budget enforcement may use stale spend values. " - "user_id=%s, response_cost=%s - %s", + "user_id=%s, response_cost=%s - %s\n%s", user_id, response_cost, str(e), + traceback.format_exc(), ) ### UPDATE END-USER SPEND ### @@ -1814,10 +1815,11 @@ async def update_cache( # noqa: PLR0915 verbose_proxy_logger.warning( "Spend tracking - failed to update end user spend in cache. " "Budget enforcement may use stale spend values. " - "end_user_id=%s, response_cost=%s - %s", + "end_user_id=%s, response_cost=%s - %s\n%s", end_user_id, response_cost, str(e), + traceback.format_exc(), ) ### UPDATE TEAM SPEND ### @@ -1861,10 +1863,11 @@ async def update_cache( # noqa: PLR0915 verbose_proxy_logger.warning( "Spend tracking - failed to update team spend in cache. " "Budget enforcement may use stale spend values. " - "team_id=%s, response_cost=%s - %s", + "team_id=%s, response_cost=%s - %s\n%s", team_id, response_cost, str(e), + traceback.format_exc(), ) ### UPDATE TAG SPEND ### @@ -1912,10 +1915,11 @@ async def update_cache( # noqa: PLR0915 verbose_proxy_logger.warning( "Spend tracking - failed to update tag spend in cache. " "Budget enforcement may use stale spend values. " - "tags=%s, response_cost=%s - %s", + "tags=%s, response_cost=%s - %s\n%s", tags, response_cost, str(e), + traceback.format_exc(), ) if token is not None and response_cost is not None: From 70ef4d0d69e3915df64c847582e5b0d60ab1fae3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 24 Feb 2026 12:10:19 -0800 Subject: [PATCH 028/132] address greptile review feedback (greploop iteration 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove re-raise in _store_transactions_in_redis so one Redis push failure doesn't drop remaining transaction types - Downgrade per-push success log from info to debug to reduce noise - Fix misleading error message in update_database — entity spend updates run as independent tasks and are not affected by this catch --- litellm/proxy/db/db_spend_update_writer.py | 3 ++- litellm/proxy/db/db_transaction_queue/redis_update_buffer.py | 3 +-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e2d365ee465..4a8c9d33d9f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -225,7 +225,8 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("Runs spend update on all tables") except Exception: verbose_proxy_logger.error( - "Spend tracking - update_database failed. All spend updates for this request will be lost. " + "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " + "may not have completed for this request. " "response_cost=%s, token=%s, user_id=%s, team_id=%s, org_id=%s, end_user_id=%s - %s", response_cost, token, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 85e139b4de0..027f3e639e3 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -101,7 +101,7 @@ class RedisUpdateBuffer: key=redis_key, values=list_of_transactions, ) - verbose_proxy_logger.info( + verbose_proxy_logger.debug( "Spend tracking - pushed spend updates to Redis buffer. " "redis_key=%s, buffer_size=%s", redis_key, @@ -118,7 +118,6 @@ class RedisUpdateBuffer: redis_key, str(e), ) - raise async def store_in_memory_spend_updates_in_redis( self, From 33719e6b38d830c5f63bf6866fd88bb996eebcca Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 12:30:18 -0800 Subject: [PATCH 029/132] docs: update v1.81.12-stable release notes to point to v1.81.12-stable.1 (#22036) --- docs/my-website/release_notes/v1.81.12.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/my-website/release_notes/v1.81.12.md b/docs/my-website/release_notes/v1.81.12.md index a1f1daa2b92..0b7c1e146ab 100644 --- a/docs/my-website/release_notes/v1.81.12.md +++ b/docs/my-website/release_notes/v1.81.12.md @@ -1,5 +1,5 @@ --- -title: "v1.81.12-stable - Guardrail Policy Templates & Action Builder" +title: "v1.81.12-stable.1 - Guardrail Policy Templates & Action Builder" slug: "v1-81-12" date: 2026-02-14T00:00:00 authors: @@ -27,7 +27,7 @@ import Image from '@theme/IdealImage'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:main-v1.81.12-stable +ghcr.io/berriai/litellm:main-v1.81.12-stable.1 ``` From 8b56e1d969385b066f596400e084e754e438fbe5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 24 Feb 2026 12:37:43 -0800 Subject: [PATCH 030/132] trigger review From 98b496433086b5a56f87e610c0349607b7553f55 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 24 Feb 2026 12:31:45 -0800 Subject: [PATCH 031/132] perf(proxy): pipeline Redis RPUSH/LPOP in spend update cycle Replace 14 sequential Redis round-trips (7 RPUSH + 7 LPOP) per spend update cycle with 2 pipelined calls (1 RPUSH pipeline + 1 LPOP pipeline). This reduces connection pool contention at scale (50+ pods). - Add RedisPipelineRpushOperation and RedisPipelineLpopOperation TypedDicts - Add async_rpush_pipeline() and async_lpop_pipeline() to RedisCache - Refactor store_in_memory_spend_updates_in_redis() to use pipeline - Add get_all_transactions_from_redis_buffer_pipeline() for batched drain - Update _commit_spend_updates_to_db_with_redis() to use pipeline drain - Existing individual methods preserved for backward compatibility --- litellm/caching/redis_cache.py | 187 ++++++++++++++- litellm/proxy/db/db_spend_update_writer.py | 31 +-- .../redis_update_buffer.py | 149 ++++++++---- litellm/types/caching.py | 18 ++ .../test_litellm/caching/test_redis_cache.py | 213 ++++++++++++++++++ .../test_redis_update_buffer.py | 194 ++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 43 ++++ 7 files changed, 773 insertions(+), 62 deletions(-) create mode 100644 tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index dcc2df5f91c..9b93d890cec 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -22,7 +22,11 @@ from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_REDIS_MAJOR_VERSION from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker -from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.caching import ( + RedisPipelineIncrementOperation, + RedisPipelineLpopOperation, + RedisPipelineRpushOperation, +) from litellm.types.services import ServiceTypes from .base_cache import BaseCache @@ -1320,6 +1324,75 @@ class RedisCache(BaseCache): ) raise e + async def _pipeline_rpush_helper( + self, + pipe: pipeline, + rpush_list: List[RedisPipelineRpushOperation], + ) -> List[int]: + """Helper function for pipeline rpush operations""" + for rpush_op in rpush_list: + pipe.rpush(rpush_op["key"], *rpush_op["values"]) + results = await pipe.execute() + # Preserve positional correspondence — raise on per-command errors + for r in results: + if isinstance(r, Exception): + raise r + return results + + async def async_rpush_pipeline( + self, + rpush_list: List[RedisPipelineRpushOperation], + ) -> List[int]: + """ + Use Redis Pipelines for bulk RPUSH operations + + Args: + rpush_list: List of RedisPipelineRpushOperation dicts containing: + - key: str + - values: List[Any] + + Returns: + List[int]: List lengths after each push + """ + if len(rpush_list) == 0: + return [] + + _redis_client: Any = self.init_async_client() + start_time = time.time() + + try: + async with _redis_client.pipeline(transaction=False) as pipe: + results = await self._pipeline_rpush_helper(pipe, rpush_list) + + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}", + ) + ) + return results + except Exception as e: + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=_duration, + error=e, + call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}", + ) + ) + verbose_logger.error( + "LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s", + str(e), + ) + raise e + async def handle_lpop_count_for_older_redis_versions( self, pipe: pipeline, key: str, count: int ) -> List[bytes]: @@ -1400,3 +1473,115 @@ class RedisCache(BaseCache): f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}" ) raise e + + async def _pipeline_lpop_helper( + self, + pipe: pipeline, + lpop_list: List[RedisPipelineLpopOperation], + ) -> List[Optional[List[str]]]: + """Helper function for pipeline lpop operations. + + For Redis >= 7, queues one LPOP(key, count) per operation. + For Redis < 7, queues `count` individual LPOP(key) commands per operation. + """ + major_version = self._parse_redis_major_version() + + if major_version >= 7: + for lpop_op in lpop_list: + pipe.lpop(lpop_op["key"], lpop_op["count"]) + raw_results = await pipe.execute() + else: + # For Redis < 7, LPOP doesn't support count param. + # Issue `count` individual LPOP commands per key, all in one pipeline. + counts: List[int] = [] + for lpop_op in lpop_list: + count = lpop_op["count"] or 1 + counts.append(count) + for _ in range(count): + pipe.lpop(lpop_op["key"]) + flat_results = await pipe.execute() + + # Re-group the flat results back into per-key lists + raw_results = [] + offset = 0 + for count in counts: + key_results = [ + r for r in flat_results[offset : offset + count] if r is not None + ] + raw_results.append(key_results if key_results else None) + offset += count + + # Decode bytes -> str for each result set + decoded_results: List[Optional[List[str]]] = [] + for r in raw_results: + if r is None: + decoded_results.append(None) + elif isinstance(r, list): + try: + decoded_results.append( + [ + item.decode("utf-8") if isinstance(item, bytes) else item + for item in r + if item is not None + ] + or None + ) + except Exception: + decoded_results.append(r) # type: ignore + else: + decoded_results.append(None) + return decoded_results + + async def async_lpop_pipeline( + self, + lpop_list: List[RedisPipelineLpopOperation], + ) -> List[Optional[List[str]]]: + """ + Use Redis Pipelines for bulk LPOP operations + + Args: + lpop_list: List of RedisPipelineLpopOperation dicts containing: + - key: str + - count: Optional[int] + + Returns: + List[Optional[List[str]]]: Decoded results per key, None if key was empty + """ + if len(lpop_list) == 0: + return [] + + _redis_client: Any = self.init_async_client() + start_time = time.time() + + try: + async with _redis_client.pipeline(transaction=False) as pipe: + results = await self._pipeline_lpop_helper(pipe, lpop_list) + + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}", + ) + ) + return results + except Exception as e: + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=_duration, + error=e, + call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}", + ) + ) + verbose_logger.error( + "LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s", + str(e), + ) + raise e diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 03628fda47f..b9cd794e780 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -509,9 +509,16 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("acquired lock for spend updates") try: - db_spend_update_transactions = ( - await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer() - ) + ( + db_spend_update_transactions, + daily_spend_update_transactions, + daily_team_spend_update_transactions, + daily_org_spend_update_transactions, + 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() + if db_spend_update_transactions is not None: await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -520,9 +527,6 @@ class DBSpendUpdateWriter: db_spend_update_transactions=db_spend_update_transactions, ) - daily_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer() - ) if daily_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_user_spend( n_retry_times=n_retry_times, @@ -530,9 +534,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_spend_update_transactions, ) - daily_team_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_team_spend_update_transactions_from_redis_buffer() - ) if daily_team_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_team_spend( n_retry_times=n_retry_times, @@ -541,9 +542,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_team_spend_update_transactions, ) - daily_org_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer() - ) if daily_org_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_org_spend( n_retry_times=n_retry_times, @@ -552,9 +550,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_org_spend_update_transactions, ) - daily_tag_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() - ) if daily_tag_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_tag_spend( n_retry_times=n_retry_times, @@ -562,9 +557,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_tag_spend_update_transactions, ) - daily_end_user_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer() - ) if daily_end_user_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_end_user_spend( n_retry_times=n_retry_times, @@ -572,9 +564,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_end_user_spend_update_transactions, ) - daily_agent_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer() - ) if daily_agent_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_agent_spend( n_retry_times=n_retry_times, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 37b42e26bc9..ef8dbb2c7b8 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -6,7 +6,7 @@ This is to prevent deadlocks and improve reliability import asyncio import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache @@ -36,6 +36,7 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.secret_managers.main import str_to_bool +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation from litellm.types.services import ServiceTypes if TYPE_CHECKING: @@ -195,47 +196,44 @@ class RedisUpdateBuffer: "ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions ) - await self._store_transactions_in_redis( - transactions=db_spend_update_transactions, - redis_key=REDIS_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_SPEND_UPDATE_QUEUE, + # Build a list of rpush operations, skipping empty/None transaction sets + _queue_configs: List[Tuple[Any, str, ServiceTypes]] = [ + (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE), + (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), + (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE), + (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), + (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE), + (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE), + (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE), + ] + + rpush_list: List[RedisPipelineRpushOperation] = [] + service_types: List[ServiceTypes] = [] + for transactions, redis_key, service_type in _queue_configs: + if transactions is None or len(transactions) == 0: + continue + rpush_list.append( + RedisPipelineRpushOperation( + key=redis_key, + values=[safe_dumps(transactions)], + ) + ) + service_types.append(service_type) + + if len(rpush_list) == 0: + return + + result_lengths = await self.redis_cache.async_rpush_pipeline( + rpush_list=rpush_list, ) - await self._store_transactions_in_redis( - transactions=daily_spend_update_transactions, - redis_key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_team_spend_update_transactions, - redis_key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_org_spend_update_transactions, - redis_key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_end_user_spend_update_transactions, - redis_key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_agent_spend_update_transactions, - redis_key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_tag_spend_update_transactions, - redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, - ) + # Emit gauge events for each queue + for i, queue_size in enumerate(result_lengths): + if i < len(service_types): + await self._emit_new_item_added_to_redis_buffer_event( + queue_size=queue_size, + service=service_types[i], + ) @staticmethod def _number_of_transactions_to_store_in_redis( @@ -317,6 +315,77 @@ class RedisUpdateBuffer: return combined_transaction + async def get_all_transactions_from_redis_buffer_pipeline( + self, + ) -> Tuple[ + Optional[DBSpendUpdateTransactions], + Optional[Dict[str, DailyUserSpendTransaction]], + Optional[Dict[str, DailyTeamSpendTransaction]], + Optional[Dict[str, DailyOrganizationSpendTransaction]], + Optional[Dict[str, DailyEndUserSpendTransaction]], + Optional[Dict[str, DailyAgentSpendTransaction]], + Optional[Dict[str, DailyTagSpendTransaction]], + ]: + """ + Drains all 7 Redis buffer queues in a single pipeline round-trip. + + Returns a 7-tuple of parsed results in this order: + 0: DBSpendUpdateTransactions + 1: daily user spend + 2: daily team spend + 3: daily org spend + 4: daily end-user spend + 5: daily agent spend + 6: daily tag spend + """ + if self.redis_cache is None: + return None, None, None, None, None, None, None + + lpop_list: List[RedisPipelineLpopOperation] = [ + RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + ] + + raw_results = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + # Pad with None if pipeline returned fewer results than expected + while len(raw_results) < 7: + raw_results.append(None) + + # Slot 0: DBSpendUpdateTransactions + db_spend: Optional[DBSpendUpdateTransactions] = None + if raw_results[0] is not None: + parsed = self._parse_list_of_transactions(raw_results[0]) + if len(parsed) > 0: + db_spend = self._combine_list_of_transactions(parsed) + + # Slots 1-6: daily spend categories + daily_results: List[Optional[Dict[str, Any]]] = [] + for slot in range(1, 7): + if raw_results[slot] is None: + daily_results.append(None) + else: + list_of_daily = [json.loads(t) for t in raw_results[slot]] # type: ignore + aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + list_of_daily + ) + daily_results.append(aggregated) + + return ( + db_spend, + cast(Optional[Dict[str, DailyUserSpendTransaction]], daily_results[0]), + cast(Optional[Dict[str, DailyTeamSpendTransaction]], daily_results[1]), + cast(Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2]), + cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]), + cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]), + cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]), + ) + async def get_all_daily_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyUserSpendTransaction]]: diff --git a/litellm/types/caching.py b/litellm/types/caching.py index ad2ffeaf9bd..7126ba3e9b9 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -52,6 +52,24 @@ class RedisPipelineSetOperation(TypedDict): ttl: Optional[int] +class RedisPipelineRpushOperation(TypedDict): + """ + TypedDict for 1 Redis Pipeline RPUSH Operation + """ + + key: str + values: List[Any] + + +class RedisPipelineLpopOperation(TypedDict): + """ + TypedDict for 1 Redis Pipeline LPOP Operation + """ + + key: str + count: Optional[int] + + DynamicCacheControl = TypedDict( "DynamicCacheControl", { diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index f3e7953b8ae..31eeb854883 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -122,6 +122,219 @@ async def test_handle_lpop_count_for_older_redis_versions(monkeypatch): assert mock_pipeline.execute.call_count == 2 +@pytest.mark.asyncio +async def test_async_rpush_pipeline_executes_all_operations(monkeypatch, redis_no_ping): + """Verify that multiple rpush ops are batched into a single pipeline execute""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.rpush = MagicMock() + mock_pipeline.execute = AsyncMock(return_value=[3, 5, 1]) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineRpushOperation + + rpush_list = [ + RedisPipelineRpushOperation(key="key1", values=["a", "b"]), + RedisPipelineRpushOperation(key="key2", values=["c"]), + RedisPipelineRpushOperation(key="key3", values=["d", "e", "f"]), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + result = await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + + assert result == [3, 5, 1] + assert mock_pipeline.rpush.call_count == 3 + mock_pipeline.rpush.assert_any_call("key1", "a", "b") + mock_pipeline.rpush.assert_any_call("key2", "c") + mock_pipeline.rpush.assert_any_call("key3", "d", "e", "f") + mock_pipeline.execute.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_rpush_pipeline_empty_list_returns_empty(monkeypatch, redis_no_ping): + """Empty rpush_list should return empty list without touching Redis""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + result = await redis_cache.async_rpush_pipeline(rpush_list=[]) + + assert result == [] + mock_redis_instance.pipeline.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_rpush_pipeline_raises_on_redis_error(monkeypatch, redis_no_ping): + """Pipeline errors should propagate""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.rpush = MagicMock() + mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down")) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineRpushOperation + + rpush_list = [RedisPipelineRpushOperation(key="key1", values=["a"])] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with pytest.raises(ConnectionError, match="Redis down"): + await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + + +@pytest.mark.asyncio +async def test_async_lpop_pipeline_single_round_trip(monkeypatch, redis_no_ping): + """Verify that multiple lpop ops are batched into a single pipeline execute""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + redis_cache.redis_version = "7.0.0" + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.lpop = MagicMock() + mock_pipeline.execute = AsyncMock(return_value=[ + [b"val1", b"val2"], # key1 results + None, # key2 empty + [b"val3"], # key3 results + ]) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineLpopOperation + + lpop_list = [ + RedisPipelineLpopOperation(key="key1", count=10), + RedisPipelineLpopOperation(key="key2", count=10), + RedisPipelineLpopOperation(key="key3", count=5), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + assert len(results) == 3 + assert results[0] == ["val1", "val2"] + assert results[1] is None + assert results[2] == ["val3"] + mock_pipeline.execute.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_lpop_pipeline_redis_lt7_regroups_flat_results(monkeypatch, redis_no_ping): + """Verify Redis < 7 fallback issues individual LPOPs and regroups correctly""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + redis_cache.redis_version = "6.2.0" + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.lpop = MagicMock() + + # With count=3 for key1 and count=2 for key2, we get 5 individual LPOP commands + # Simulate: key1 has 2 values then None, key2 has 1 value then None + mock_pipeline.execute = AsyncMock(return_value=[ + b"val1", b"val2", None, # 3 LPOPs for key1 + b"val3", None, # 2 LPOPs for key2 + ]) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineLpopOperation + + lpop_list = [ + RedisPipelineLpopOperation(key="key1", count=3), + RedisPipelineLpopOperation(key="key2", count=2), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + assert len(results) == 2 + assert results[0] == ["val1", "val2"] # 2 values, None filtered out + assert results[1] == ["val3"] # 1 value, None filtered out + # All 5 individual LPOPs should be queued, but only 1 execute() call + assert mock_pipeline.lpop.call_count == 5 + mock_pipeline.execute.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_rpush_pipeline_raises_on_per_command_error(monkeypatch, redis_no_ping): + """Verify that per-command errors in pipeline results are raised, not silently dropped""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.rpush = MagicMock() + # Simulate: first RPUSH succeeds, second returns a per-command error + mock_pipeline.execute = AsyncMock(return_value=[3, Exception("WRONGTYPE")]) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineRpushOperation + + rpush_list = [ + RedisPipelineRpushOperation(key="key1", values=["a"]), + RedisPipelineRpushOperation(key="key2", values=["b"]), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with pytest.raises(Exception, match="WRONGTYPE"): + await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + + +@pytest.mark.asyncio +async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping): + """Empty lpop_list should return empty list without touching Redis""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + result = await redis_cache.async_lpop_pipeline(lpop_list=[]) + + assert result == [] + mock_redis_instance.pipeline.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_lpop_pipeline_propagates_redis_exception(monkeypatch, redis_no_ping): + """Pipeline errors should propagate""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + redis_cache.redis_version = "7.0.0" + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.lpop = MagicMock() + mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down")) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineLpopOperation + + lpop_list = [RedisPipelineLpopOperation(key="key1", count=10)] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with pytest.raises(ConnectionError, match="Redis down"): + await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + @pytest.mark.asyncio @pytest.mark.parametrize( "redis_version", diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py new file mode 100644 index 00000000000..2a380370c30 --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -0,0 +1,194 @@ +import json +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 + +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.types.caching import RedisPipelineRpushOperation + + +@pytest.fixture +def mock_redis_cache(): + """Create a mock RedisCache instance""" + mock = AsyncMock() + return mock + + +@pytest.fixture +def redis_update_buffer(mock_redis_cache): + """Create a RedisUpdateBuffer with a mock RedisCache""" + return RedisUpdateBuffer(redis_cache=mock_redis_cache) + + +@pytest.mark.asyncio +async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache): + """ + Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once + with the correct operations and skips empty queues. + """ + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[3, 5, 2]) + + # Create mock queues - only 3 of 7 have data + spend_update_queue = AsyncMock() + spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( + return_value={"key_list_transactions": {"key1": 1.0}} + ) + + daily_spend_queue = AsyncMock() + daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"user_key1": {"spend": 1.0}} + ) + + daily_team_queue = AsyncMock() + daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"team_key1": {"spend": 2.0}} + ) + + # Empty queues + daily_org_queue = AsyncMock() + daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={} + ) + + daily_end_user_queue = AsyncMock() + daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value=None + ) + + daily_agent_queue = AsyncMock() + daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={} + ) + + daily_tag_queue = AsyncMock() + daily_tag_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={} + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=spend_update_queue, + daily_spend_update_queue=daily_spend_queue, + daily_team_spend_update_queue=daily_team_queue, + daily_org_spend_update_queue=daily_org_queue, + daily_end_user_spend_update_queue=daily_end_user_queue, + daily_agent_spend_update_queue=daily_agent_queue, + daily_tag_spend_update_queue=daily_tag_queue, + ) + + # Should be called exactly once (pipeline) + mock_redis_cache.async_rpush_pipeline.assert_called_once() + + # Verify only 3 operations were included (empty ones skipped) + call_args = mock_redis_cache.async_rpush_pipeline.call_args + rpush_list = call_args.kwargs["rpush_list"] + assert len(rpush_list) == 3 + + +@pytest.mark.asyncio +async def test_store_in_memory_spend_updates_all_empty_returns_early( + redis_update_buffer, mock_redis_cache +): + """ + When all queues are empty, pipeline should never be called. + """ + mock_redis_cache.async_rpush_pipeline = AsyncMock() + + # All queues return empty + empty_queue = AsyncMock() + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( + return_value={} + ) + empty_daily_queue = AsyncMock() + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={} + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=empty_queue, + daily_spend_update_queue=empty_daily_queue, + daily_team_spend_update_queue=empty_daily_queue, + daily_org_spend_update_queue=empty_daily_queue, + daily_end_user_spend_update_queue=empty_daily_queue, + daily_agent_spend_update_queue=empty_daily_queue, + daily_tag_spend_update_queue=empty_daily_queue, + ) + + mock_redis_cache.async_rpush_pipeline.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_all_transactions_from_redis_buffer_pipeline( + redis_update_buffer, mock_redis_cache +): + """ + Verify get_all_transactions_from_redis_buffer_pipeline correctly parses + and aggregates results from async_lpop_pipeline. + """ + # Simulate pipeline results: slot 0 = spend updates, slots 1-6 = daily categories + db_spend_json = json.dumps( + { + "key_list_transactions": {"key1": 1.0, "key2": 2.0}, + "user_list_transactions": {"user1": 0.5}, + "end_user_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + } + ) + daily_user_json = json.dumps({"user_key1": {"spend": 1.0, "api_requests": 1}}) + daily_team_json = json.dumps({"team_key1": {"spend": 2.0, "api_requests": 2}}) + + mock_redis_cache.async_lpop_pipeline = AsyncMock( + return_value=[ + [db_spend_json], # slot 0: db spend updates + [daily_user_json], # slot 1: daily user + [daily_team_json], # slot 2: daily team + None, # slot 3: daily org (empty) + None, # slot 4: daily end-user (empty) + None, # slot 5: daily agent (empty) + None, # slot 6: daily tag (empty) + ] + ) + + result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + + assert len(result) == 7 + db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent, daily_tag = result + + # Verify db spend was parsed correctly + assert db_spend is not None + assert db_spend["key_list_transactions"]["key1"] == 1.0 + assert db_spend["key_list_transactions"]["key2"] == 2.0 + assert db_spend["user_list_transactions"]["user1"] == 0.5 + + # Verify daily user was parsed + assert daily_user is not None + assert daily_user["user_key1"]["spend"] == 1.0 + + # Verify daily team was parsed + assert daily_team is not None + assert daily_team["team_key1"]["spend"] == 2.0 + + # Verify empty slots + assert daily_org is None + assert daily_end_user is None + assert daily_agent is None + assert daily_tag is None + + # Verify pipeline was called once with correct keys + mock_redis_cache.async_lpop_pipeline.assert_called_once() + + +@pytest.mark.asyncio +async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): + """When redis_cache is None, should return all Nones""" + buffer = RedisUpdateBuffer(redis_cache=None) + result = await buffer.get_all_transactions_from_redis_buffer_pipeline() + assert result == (None, None, None, None, None, None, None) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 0fa0d4cf10b..2a0428e2f07 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1076,3 +1076,46 @@ async def test_commit_key_spend_updates_includes_last_active(): last_active = call_kwargs["data"]["last_active"] assert isinstance(last_active, datetime) assert before_call <= last_active <= after_call + + +@pytest.mark.asyncio +async def test_commit_spend_updates_uses_pipeline(): + """ + Verify that _commit_spend_updates_to_db_with_redis uses + get_all_transactions_from_redis_buffer_pipeline instead of 7 individual calls. + """ + db_writer = DBSpendUpdateWriter() + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() + # Return all-None tuple (no data to commit) + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) + ) + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + mock_prisma_client = MagicMock() + mock_proxy_logging = MagicMock() + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=mock_prisma_client, + n_retry_times=1, + proxy_logging_obj=mock_proxy_logging, + ) + + # Pipeline method should be called once + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline.assert_called_once() + + # Individual methods should NOT be called + mock_redis_update_buffer.get_all_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_team_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() From 2cabbccf6f8244348dc908e2a2d0a069b3313001 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 24 Feb 2026 12:49:48 -0800 Subject: [PATCH 032/132] address greptile review feedback (greploop iteration 3) - Add missing traceback to team member spend enqueue error log --- litellm/proxy/db/db_spend_update_writer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4a8c9d33d9f..d7d6b2b1eb0 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -351,11 +351,12 @@ class DBSpendUpdateWriter: except Exception as e: verbose_proxy_logger.error( "Spend tracking - failed to enqueue team member spend update. " - "team_id=%s, user_id=%s, response_cost=%s - %s", + "team_id=%s, user_id=%s, response_cost=%s - %s\n%s", team_id, user_id, response_cost, str(e), + traceback.format_exc(), ) except Exception as e: verbose_proxy_logger.error( From 68a30a39e68d631959920f7c658baaf8bea2df06 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 24 Feb 2026 14:22:57 -0800 Subject: [PATCH 033/132] fix(proxy): add LPOP pipeline error checking and fix org spend ServiceType - Add per-command error check in _pipeline_lpop_helper to match _pipeline_rpush_helper, preventing silent data loss on WRONGTYPE errors - Fix pre-existing bug: org spend queue metric was using REDIS_DAILY_SPEND_UPDATE_QUEUE instead of REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE - Add test for per-command LPOP pipeline error propagation --- litellm/caching/redis_cache.py | 5 ++++ .../redis_update_buffer.py | 2 +- .../test_litellm/caching/test_redis_cache.py | 30 +++++++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 9b93d890cec..fa9b94bc2ac 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1511,6 +1511,11 @@ class RedisCache(BaseCache): raw_results.append(key_results if key_results else None) offset += count + # Raise on per-command errors (matches _pipeline_rpush_helper behavior) + for r in raw_results: + if isinstance(r, Exception): + raise r + # Decode bytes -> str for each result set decoded_results: List[Optional[List[str]]] = [] for r in raw_results: diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index ef8dbb2c7b8..504fc5aba02 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -201,7 +201,7 @@ class RedisUpdateBuffer: (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE), (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE), - (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), + (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE), (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE), diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 31eeb854883..82606511826 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -296,6 +296,36 @@ async def test_async_rpush_pipeline_raises_on_per_command_error(monkeypatch, red await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) +@pytest.mark.asyncio +async def test_async_lpop_pipeline_raises_on_per_command_error(monkeypatch, redis_no_ping): + """Verify that per-command errors in LPOP pipeline results are raised, not silently dropped""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + redis_cache.redis_version = "7.0.0" + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.lpop = MagicMock() + # Simulate: first LPOP succeeds, second returns a per-command error + mock_pipeline.execute = AsyncMock( + return_value=[[b"val1"], Exception("WRONGTYPE")] + ) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineLpopOperation + + lpop_list = [ + RedisPipelineLpopOperation(key="key1", count=10), + RedisPipelineLpopOperation(key="key2", count=10), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with pytest.raises(Exception, match="WRONGTYPE"): + await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + @pytest.mark.asyncio async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping): """Empty lpop_list should return empty list without touching Redis""" From 9971b67587f8c376923995ee1260429195458124 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 24 Feb 2026 15:21:43 -0800 Subject: [PATCH 034/132] fix: convert remaining dict(request.headers) to _safe_get_request_headers Missed conversions in user_api_key_auth.py and litellm_pre_call_utils.py. Both call sites are read-only so no .copy() needed. --- litellm/proxy/auth/user_api_key_auth.py | 2 +- litellm/proxy/litellm_pre_call_utils.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ca9af16ff91..8f17440773a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -483,7 +483,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 parent_otel_span = ( open_telemetry_logger.create_litellm_proxy_request_started_span( start_time=start_time, - headers=dict(request.headers), + headers=_safe_get_request_headers(request), ) ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b61dfa5b263..52f0b1d46e9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -14,6 +14,7 @@ from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors, LitellmDataForBackendLLMCall, LitellmUserRoles, SpecialHeaders, TeamCallbackMetadata, UserAPIKeyAuth) +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers # Cache special headers as a frozenset for O(1) lookup performance _SPECIAL_HEADERS_CACHE = frozenset( @@ -824,7 +825,7 @@ 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 SecretFields - _raw_headers: Dict[str, str] = dict(request.headers) + _raw_headers: Dict[str, str] = _safe_get_request_headers(request) _headers: Dict[str, str] = clean_headers( request.headers, litellm_key_header_name=( From b3bb744aa4d0fa026e93f89ad6a11ce75b86dbdb Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 24 Feb 2026 15:26:06 -0800 Subject: [PATCH 035/132] [Test] Add unit tests for router_settings components Co-Authored-By: Claude Sonnet 4.6 --- .../LatencyBasedConfiguration.test.tsx | 55 ++++++ .../ReliabilityRetriesSection.test.tsx | 80 +++++++++ .../RouterSettingsForm.test.tsx | 134 ++++++++++++++ .../RoutingStrategySelector.test.tsx | 98 ++++++++++ .../TagFilteringToggle.test.tsx | 113 ++++++++++++ .../components/router_settings/index.test.tsx | 170 ++++++++++++++++++ 6 files changed, 650 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx create mode 100644 ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx create mode 100644 ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx create mode 100644 ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx create mode 100644 ui/litellm-dashboard/src/components/router_settings/index.test.tsx diff --git a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx new file mode 100644 index 00000000000..0176be5ab40 --- /dev/null +++ b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx @@ -0,0 +1,55 @@ +import { describe, it, expect } from "vitest"; +import { render, screen } from "@testing-library/react"; +import LatencyBasedConfiguration from "./LatencyBasedConfiguration"; + +describe("LatencyBasedConfiguration", () => { + it("should render the section heading", () => { + render(); + expect(screen.getByText("Latency-Based Configuration")).toBeInTheDocument(); + }); + + it("should render default params when no args are provided", () => { + render(); + // Default: ttl=3600, lowest_latency_buffer=0 + expect(screen.getByDisplayValue("3600")).toBeInTheDocument(); + expect(screen.getByDisplayValue("0")).toBeInTheDocument(); + }); + + it("should render the provided routing strategy args as inputs", () => { + const args = { ttl: 7200, lowest_latency_buffer: 0.1 }; + render(); + expect(screen.getByDisplayValue("7200")).toBeInTheDocument(); + expect(screen.getByDisplayValue("0.1")).toBeInTheDocument(); + }); + + it("should render an input with the correct name attribute for each param", () => { + const args = { ttl: 3600, lowest_latency_buffer: 0 }; + render(); + expect(screen.getByRole("textbox", { name: /ttl/i })).toBeInTheDocument(); + expect(screen.getByRole("textbox", { name: /lowest latency buffer/i })).toBeInTheDocument(); + }); + + it("should display the TTL parameter explanation", () => { + render(); + expect( + screen.getByText(/sliding window to look back over/i) + ).toBeInTheDocument(); + }); + + it("should display the lowest_latency_buffer parameter explanation", () => { + render(); + expect( + screen.getByText(/shuffle between deployments within this %/i) + ).toBeInTheDocument(); + }); + + it("should render object values stringified into the input", () => { + const args = { ttl: { nested: true } }; + render(); + // HTML input type=text strips newlines, so check that the key/value appears + const input = document.querySelector('input[name="ttl"]') as HTMLInputElement; + expect(input).not.toBeNull(); + expect(input.value).toContain('"nested"'); + expect(input.value).toContain('true'); + }); +}); diff --git a/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx b/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx new file mode 100644 index 00000000000..0892d39fc29 --- /dev/null +++ b/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx @@ -0,0 +1,80 @@ +import { describe, it, expect } from "vitest"; +import { render, screen } from "@testing-library/react"; +import ReliabilityRetriesSection from "./ReliabilityRetriesSection"; + +const baseSettings = { + num_retries: 3, + timeout: 30, + allowed_fails: 2, + fallbacks: ["gpt-3.5"], + context_window_fallbacks: [], + routing_strategy_args: { ttl: 3600 }, + routing_strategy: "simple-shuffle", + enable_tag_filtering: false, +}; + +describe("ReliabilityRetriesSection", () => { + it("should render the section heading", () => { + render(); + expect(screen.getByText("Reliability & Retries")).toBeInTheDocument(); + }); + + it("should render input fields for non-excluded settings", () => { + render(); + expect(screen.getByDisplayValue("3")).toBeInTheDocument(); // num_retries + expect(screen.getByDisplayValue("30")).toBeInTheDocument(); // timeout + expect(screen.getByDisplayValue("2")).toBeInTheDocument(); // allowed_fails + }); + + it("should not render inputs for excluded keys", () => { + render(); + // Each excluded key must not produce a visible input value + const inputs = screen.queryAllByRole("textbox"); + const inputNames = inputs.map((el) => el.getAttribute("name")); + expect(inputNames).not.toContain("fallbacks"); + expect(inputNames).not.toContain("context_window_fallbacks"); + expect(inputNames).not.toContain("routing_strategy_args"); + expect(inputNames).not.toContain("routing_strategy"); + expect(inputNames).not.toContain("enable_tag_filtering"); + }); + + it("should use ui_field_name from metadata as the label", () => { + const metadata = { + num_retries: { ui_field_name: "Number of Retries", field_description: "How many times to retry" }, + }; + render( + + ); + expect(screen.getByText("Number of Retries")).toBeInTheDocument(); + }); + + it("should fall back to the raw param name when no metadata label is available", () => { + render( + + ); + expect(screen.getByText("num_retries")).toBeInTheDocument(); + }); + + it("should render null values as an empty input", () => { + render( + + ); + const input = screen.getByRole("textbox", { name: /timeout/i }) as HTMLInputElement; + expect(input.value).toBe(""); + }); + + it("should render object values stringified into the input", () => { + const settings = { retry_policy: { "rate-limited": 2 } }; + render(); + // HTML input type=text strips newlines, so check that the key/value appears + const input = document.querySelector('input[name="retry_policy"]') as HTMLInputElement; + expect(input).not.toBeNull(); + expect(input.value).toContain('"rate-limited"'); + expect(input.value).toContain('2'); + }); + + it("should render no inputs when routerSettings is empty", () => { + render(); + expect(screen.queryAllByRole("textbox")).toHaveLength(0); + }); +}); diff --git a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx new file mode 100644 index 00000000000..767820cd485 --- /dev/null +++ b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx @@ -0,0 +1,134 @@ +import { describe, it, expect, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import RouterSettingsForm from "./RouterSettingsForm"; +import type { RouterSettingsFormValue } from "./RouterSettingsForm"; + +// Use the same antd mock as RoutingStrategySelector to keep things consistent +vi.mock("antd", () => ({ + Select: Object.assign( + ({ value, onChange, children }: any) => ( + + ), + { + Option: ({ value, children }: any) => ( + + ), + } + ), +})); + +vi.mock("@tremor/react", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + Switch: ({ checked, onChange }: any) => ( + onChange(e.target.checked)} + /> + ), + }; +}); + +const defaultValue: RouterSettingsFormValue = { + routerSettings: {}, + selectedStrategy: null, + enableTagFiltering: false, +}; + +const baseProps = { + value: defaultValue, + onChange: vi.fn(), + routerFieldsMetadata: {}, + availableRoutingStrategies: [], + routingStrategyDescriptions: {}, +}; + +describe("RouterSettingsForm", () => { + it("should render", () => { + render(); + expect(screen.getByText("Routing Settings")).toBeInTheDocument(); + }); + + it("should not show the strategy selector when no strategies are provided", () => { + render(); + expect(screen.queryByTestId("strategy-select")).not.toBeInTheDocument(); + }); + + it("should show the strategy selector when strategies are available", () => { + const props = { + ...baseProps, + availableRoutingStrategies: ["simple-shuffle", "latency-based-routing"], + }; + render(); + expect(screen.getByTestId("strategy-select")).toBeInTheDocument(); + }); + + it("should not render LatencyBasedConfiguration for non-latency strategies", () => { + const props = { + ...baseProps, + value: { ...defaultValue, selectedStrategy: "simple-shuffle" }, + availableRoutingStrategies: ["simple-shuffle"], + }; + render(); + expect(screen.queryByText("Latency-Based Configuration")).not.toBeInTheDocument(); + }); + + it("should render LatencyBasedConfiguration when strategy is latency-based-routing", () => { + const props = { + ...baseProps, + value: { + ...defaultValue, + selectedStrategy: "latency-based-routing", + routerSettings: { routing_strategy_args: { ttl: 3600, lowest_latency_buffer: 0 } }, + }, + availableRoutingStrategies: ["latency-based-routing"], + }; + render(); + expect(screen.getByText("Latency-Based Configuration")).toBeInTheDocument(); + }); + + it("should call onChange with the updated strategy when the selector changes", () => { + const onChange = vi.fn(); + const props = { + ...baseProps, + onChange, + availableRoutingStrategies: ["simple-shuffle", "latency-based-routing"], + }; + render(); + + const select = screen.getByTestId("strategy-select") as HTMLSelectElement; + select.value = "latency-based-routing"; + select.dispatchEvent(new Event("change", { bubbles: true })); + + expect(onChange).toHaveBeenCalledWith( + expect.objectContaining({ selectedStrategy: "latency-based-routing" }) + ); + }); + + it("should call onChange with the updated enableTagFiltering when the toggle changes", async () => { + const onChange = vi.fn(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("switch")); + + expect(onChange).toHaveBeenCalledWith( + expect.objectContaining({ enableTagFiltering: true }) + ); + }); + + it("should show the Reliability & Retries section", () => { + render(); + expect(screen.getByText("Reliability & Retries")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx b/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx new file mode 100644 index 00000000000..85b1dc21acf --- /dev/null +++ b/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx @@ -0,0 +1,98 @@ +import { describe, it, expect, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import RoutingStrategySelector from "./RoutingStrategySelector"; + +// Ant Design's Select is complex to drive in JSDOM; swap it for a plain +// onChange(e.target.value)} + > + {children} + +
+ ), + { + Option: ({ value, children }: any) => ( + + ), + } + ), +})); + +const baseProps = { + selectedStrategy: null, + availableStrategies: ["simple-shuffle", "latency-based-routing", "least-busy"], + routingStrategyDescriptions: { + "simple-shuffle": "Randomly pick a deployment", + "latency-based-routing": "Pick the lowest-latency deployment", + }, + routerFieldsMetadata: {}, + onStrategyChange: vi.fn(), +}; + +describe("RoutingStrategySelector", () => { + it("should render", () => { + render(); + expect(screen.getByTestId("ant-select")).toBeInTheDocument(); + }); + + it("should display default label when no metadata is provided", () => { + render(); + expect(screen.getByText("Routing Strategy")).toBeInTheDocument(); + }); + + it("should display ui_field_name from metadata when provided", () => { + const props = { + ...baseProps, + routerFieldsMetadata: { + routing_strategy: { + ui_field_name: "Strategy", + field_description: "How to pick a deployment", + }, + }, + }; + render(); + expect(screen.getByText("Strategy")).toBeInTheDocument(); + expect(screen.getByText("How to pick a deployment")).toBeInTheDocument(); + }); + + it("should render all available strategies as options", () => { + render(); + expect(screen.getByText("simple-shuffle")).toBeInTheDocument(); + expect(screen.getByText("latency-based-routing")).toBeInTheDocument(); + expect(screen.getByText("least-busy")).toBeInTheDocument(); + }); + + it("should display strategy descriptions alongside option labels", () => { + render(); + expect(screen.getByText("Randomly pick a deployment")).toBeInTheDocument(); + expect(screen.getByText("Pick the lowest-latency deployment")).toBeInTheDocument(); + }); + + it("should not render a description for a strategy that has none", () => { + render(); + // "least-busy" has no entry in routingStrategyDescriptions + const select = screen.getByTestId("strategy-select"); + const leastBusyOption = Array.from(select.querySelectorAll("option")).find( + (o) => o.value === "least-busy" + ); + expect(leastBusyOption).toBeInTheDocument(); + }); + + it("should call onStrategyChange with the selected strategy value", () => { + const onStrategyChange = vi.fn(); + render(); + + const select = screen.getByTestId("strategy-select") as HTMLSelectElement; + select.value = "latency-based-routing"; + select.dispatchEvent(new Event("change", { bubbles: true })); + + expect(onStrategyChange).toHaveBeenCalledWith("latency-based-routing"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx b/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx new file mode 100644 index 00000000000..593721db023 --- /dev/null +++ b/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx @@ -0,0 +1,113 @@ +import { describe, it, expect, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import TagFilteringToggle from "./TagFilteringToggle"; + +// setupTests.ts mocks @tremor/react but leaves Switch as the real implementation. +// Re-mock Switch as a plain checkbox so toggle interactions are trivially testable. +vi.mock("@tremor/react", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + Switch: ({ checked, onChange, className }: any) => ( + onChange(e.target.checked)} + className={className} + /> + ), + }; +}); + +const baseMetadata = { + enable_tag_filtering: { + ui_field_name: "Tag Filtering", + field_description: "Route requests based on tags", + link: null, + }, +}; + +describe("TagFilteringToggle", () => { + it("should render", () => { + render( + + ); + expect(screen.getByRole("switch")).toBeInTheDocument(); + }); + + it("should display default label when no metadata is provided", () => { + render( + + ); + expect(screen.getByText("Enable Tag Filtering")).toBeInTheDocument(); + }); + + it("should display the label from metadata when provided", () => { + render( + + ); + expect(screen.getByText("Tag Filtering")).toBeInTheDocument(); + }); + + it("should display the description from metadata", () => { + render( + + ); + expect(screen.getByText("Route requests based on tags")).toBeInTheDocument(); + }); + + it("should render a Learn more link when metadata provides one", () => { + const metadata = { + enable_tag_filtering: { + ...baseMetadata.enable_tag_filtering, + link: "https://docs.example.com/tag-filtering", + }, + }; + render( + + ); + const link = screen.getByRole("link", { name: /learn more/i }); + expect(link).toBeInTheDocument(); + expect(link).toHaveAttribute("href", "https://docs.example.com/tag-filtering"); + }); + + it("should not render a Learn more link when metadata has no link", () => { + render( + + ); + expect(screen.queryByRole("link", { name: /learn more/i })).not.toBeInTheDocument(); + }); + + it("should reflect the enabled=true state on the switch", () => { + render( + + ); + expect(screen.getByRole("switch")).toBeChecked(); + }); + + it("should call onToggle with the new value when the switch is toggled", async () => { + const onToggle = vi.fn(); + const user = userEvent.setup(); + render( + + ); + + await user.click(screen.getByRole("switch")); + + expect(onToggle).toHaveBeenCalledWith(true); + }); +}); diff --git a/ui/litellm-dashboard/src/components/router_settings/index.test.tsx b/ui/litellm-dashboard/src/components/router_settings/index.test.tsx new file mode 100644 index 00000000000..1920268207d --- /dev/null +++ b/ui/litellm-dashboard/src/components/router_settings/index.test.tsx @@ -0,0 +1,170 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import RouterSettings from "./index"; + +vi.mock("antd", () => ({ + Select: Object.assign( + ({ value, onChange, children }: any) => ( + + ), + { + Option: ({ value, children }: any) => ( + + ), + } + ), +})); + +vi.mock("@tremor/react", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + Switch: ({ checked, onChange }: any) => ( + onChange(e.target.checked)} + /> + ), + }; +}); + +vi.mock("@/components/networking", () => ({ + getCallbacksCall: vi.fn(), + getRouterSettingsCall: vi.fn(), + setCallbacksCall: vi.fn(), +})); + +import { + getCallbacksCall, + getRouterSettingsCall, + setCallbacksCall, +} from "@/components/networking"; + +const mockCallbacksResponse = { + router_settings: { + routing_strategy: "simple-shuffle", + num_retries: 3, + timeout: 30, + }, +}; + +const mockRouterSettingsResponse = { + fields: [ + { + field_name: "routing_strategy", + ui_field_name: "Routing Strategy", + field_description: "How requests are distributed", + options: ["simple-shuffle", "latency-based-routing"], + link: null, + }, + { + field_name: "enable_tag_filtering", + ui_field_name: "Tag Filtering", + field_description: "Route by tag", + field_value: false, + link: null, + }, + ], + routing_strategy_descriptions: { + "simple-shuffle": "Randomly pick a deployment", + "latency-based-routing": "Pick the lowest-latency deployment", + }, +}; + +const defaultProps = { + accessToken: "test-token", + userRole: "Admin", + userID: "user-1", + modelData: null, +}; + +describe("RouterSettings", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(getCallbacksCall).mockResolvedValue(mockCallbacksResponse); + vi.mocked(getRouterSettingsCall).mockResolvedValue(mockRouterSettingsResponse); + vi.mocked(setCallbacksCall).mockResolvedValue({}); + }); + + it("should render nothing when accessToken is null", () => { + const { container } = renderWithProviders( + + ); + expect(container).toBeEmptyDOMElement(); + }); + + it("should render the Save Changes and Reset buttons when authenticated", () => { + renderWithProviders(); + expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /reset/i })).toBeInTheDocument(); + }); + + it("should fetch callbacks and router settings on mount", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(getCallbacksCall).toHaveBeenCalledWith("test-token", "user-1", "Admin"); + }); + expect(getRouterSettingsCall).toHaveBeenCalledWith("test-token"); + }); + + it("should not fetch data when any required prop is missing", () => { + renderWithProviders( + + ); + expect(getCallbacksCall).not.toHaveBeenCalled(); + }); + + it("should render routing strategies loaded from the API", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByTestId("strategy-select")).toBeInTheDocument(); + }); + + const select = screen.getByTestId("strategy-select") as HTMLSelectElement; + const optionValues = Array.from(select.options).map((o) => o.value); + expect(optionValues).toContain("simple-shuffle"); + expect(optionValues).toContain("latency-based-routing"); + }); + + it("should call setCallbacksCall with updated settings on Save Changes", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect(setCallbacksCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ router_settings: expect.any(Object) }) + ); + }); + + it("should show a success notification after saving", async () => { + const NotificationsManager = await import("@/components/molecules/notifications_manager"); + const user = userEvent.setup(); + renderWithProviders(); + + await waitFor(() => + expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument() + ); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect(NotificationsManager.default.success).toHaveBeenCalledWith( + "router settings updated successfully" + ); + }); +}); From 784af16cb4064fcd27f407a4b8e268045652ba62 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 24 Feb 2026 15:48:31 -0800 Subject: [PATCH 036/132] address greptile review feedback (greploop iteration 1) - Replace document.querySelector/querySelectorAll with screen.getByRole - Replace raw dispatchEvent with userEvent.selectOptions --- .../LatencyBasedConfiguration.test.tsx | 3 +-- .../ReliabilityRetriesSection.test.tsx | 3 +-- .../router_settings/RouterSettingsForm.test.tsx | 7 +++---- .../RoutingStrategySelector.test.tsx | 16 ++++++---------- 4 files changed, 11 insertions(+), 18 deletions(-) diff --git a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx index 0176be5ab40..39d9f8c881c 100644 --- a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx @@ -47,8 +47,7 @@ describe("LatencyBasedConfiguration", () => { const args = { ttl: { nested: true } }; render(); // HTML input type=text strips newlines, so check that the key/value appears - const input = document.querySelector('input[name="ttl"]') as HTMLInputElement; - expect(input).not.toBeNull(); + const input = screen.getByRole("textbox", { name: /ttl/i }) as HTMLInputElement; expect(input.value).toContain('"nested"'); expect(input.value).toContain('true'); }); diff --git a/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx b/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx index 0892d39fc29..101b09af0dc 100644 --- a/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx @@ -67,8 +67,7 @@ describe("ReliabilityRetriesSection", () => { const settings = { retry_policy: { "rate-limited": 2 } }; render(); // HTML input type=text strips newlines, so check that the key/value appears - const input = document.querySelector('input[name="retry_policy"]') as HTMLInputElement; - expect(input).not.toBeNull(); + const input = screen.getByRole("textbox", { name: /retry_policy/i }) as HTMLInputElement; expect(input.value).toContain('"rate-limited"'); expect(input.value).toContain('2'); }); diff --git a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx index 767820cd485..01f318b909a 100644 --- a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx @@ -97,8 +97,9 @@ describe("RouterSettingsForm", () => { expect(screen.getByText("Latency-Based Configuration")).toBeInTheDocument(); }); - it("should call onChange with the updated strategy when the selector changes", () => { + it("should call onChange with the updated strategy when the selector changes", async () => { const onChange = vi.fn(); + const user = userEvent.setup(); const props = { ...baseProps, onChange, @@ -106,9 +107,7 @@ describe("RouterSettingsForm", () => { }; render(); - const select = screen.getByTestId("strategy-select") as HTMLSelectElement; - select.value = "latency-based-routing"; - select.dispatchEvent(new Event("change", { bubbles: true })); + await user.selectOptions(screen.getByTestId("strategy-select"), "latency-based-routing"); expect(onChange).toHaveBeenCalledWith( expect.objectContaining({ selectedStrategy: "latency-based-routing" }) diff --git a/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx b/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx index 85b1dc21acf..01f839681f1 100644 --- a/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx @@ -1,5 +1,6 @@ import { describe, it, expect, vi } from "vitest"; import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import RoutingStrategySelector from "./RoutingStrategySelector"; // Ant Design's Select is complex to drive in JSDOM; swap it for a plain @@ -77,21 +78,16 @@ describe("RoutingStrategySelector", () => { it("should not render a description for a strategy that has none", () => { render(); - // "least-busy" has no entry in routingStrategyDescriptions - const select = screen.getByTestId("strategy-select"); - const leastBusyOption = Array.from(select.querySelectorAll("option")).find( - (o) => o.value === "least-busy" - ); - expect(leastBusyOption).toBeInTheDocument(); + // "least-busy" has no entry in routingStrategyDescriptions — it still renders without crashing + expect(screen.getByText("least-busy")).toBeInTheDocument(); }); - it("should call onStrategyChange with the selected strategy value", () => { + it("should call onStrategyChange with the selected strategy value", async () => { const onStrategyChange = vi.fn(); + const user = userEvent.setup(); render(); - const select = screen.getByTestId("strategy-select") as HTMLSelectElement; - select.value = "latency-based-routing"; - select.dispatchEvent(new Event("change", { bubbles: true })); + await user.selectOptions(screen.getByTestId("strategy-select"), "latency-based-routing"); expect(onStrategyChange).toHaveBeenCalledWith("latency-based-routing"); }); From 26e5482abb352b16dd57444705a6e44ef206d3d8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 24 Feb 2026 15:52:15 -0800 Subject: [PATCH 037/132] address greptile review feedback (greploop iteration 2) - Wait for strategy select (API data loaded) before clicking Save - Assert specific payload content in setCallbacksCall - Move NotificationsManager import to top of file --- .../components/router_settings/index.test.tsx | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/components/router_settings/index.test.tsx b/ui/litellm-dashboard/src/components/router_settings/index.test.tsx index 1920268207d..8edb0ac6e07 100644 --- a/ui/litellm-dashboard/src/components/router_settings/index.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/index.test.tsx @@ -48,6 +48,7 @@ import { getRouterSettingsCall, setCallbacksCall, } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; const mockCallbacksResponse = { router_settings: { @@ -141,29 +142,34 @@ describe("RouterSettings", () => { const user = userEvent.setup(); renderWithProviders(); + // Wait for the strategy select to appear — it only renders after getRouterSettingsCall resolves await waitFor(() => { - expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument(); + expect(screen.getByTestId("strategy-select")).toBeInTheDocument(); }); await user.click(screen.getByRole("button", { name: /save changes/i })); expect(setCallbacksCall).toHaveBeenCalledWith( "test-token", - expect.objectContaining({ router_settings: expect.any(Object) }) + expect.objectContaining({ + router_settings: expect.objectContaining({ + routing_strategy: "simple-shuffle", + }), + }) ); }); it("should show a success notification after saving", async () => { - const NotificationsManager = await import("@/components/molecules/notifications_manager"); const user = userEvent.setup(); renderWithProviders(); - await waitFor(() => - expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument() - ); + // Wait for data to load before interacting + await waitFor(() => { + expect(screen.getByTestId("strategy-select")).toBeInTheDocument(); + }); await user.click(screen.getByRole("button", { name: /save changes/i })); - expect(NotificationsManager.default.success).toHaveBeenCalledWith( + expect(NotificationsManager.success).toHaveBeenCalledWith( "router settings updated successfully" ); }); From 6ee50ff73e47da30c31526c22c8b58979248348c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 16:27:06 -0800 Subject: [PATCH 038/132] feat(proxy): tool policies - auto-discover tools + policy enforcement guardrail (#22041) * feat(proxy): tool policies - auto-discover tools, manage policies, guardrail enforcement - New LiteLLM_ToolTable in schema.prisma to store discovered tools - Auto-discovery: tools seen in LLM responses get upserted via ToolDiscoveryQueue (hooks into DBSpendUpdateWriter, same pipeline as spend tracking) - Management endpoints: GET /v1/tool/list, GET /v1/tool/{name}, POST /v1/tool/policy - ToolPolicyGuardrail: blocks tool_calls in responses based on policy setting - UI: Tool Policies page under Guardrails section with policy selector, filters by policy/team/key, live tail, sortable table - Unit tests for queue, writer, endpoints, guardrail * feat(tool-policies): track call_count + discover tools from request body and /messages API - Add call_count column to LiteLLM_ToolTable; incremented on every flush - Extract tools from request body too (not just response tool_calls): - OpenAI /chat/completions: tools[].function.name - Anthropic /messages pass-through: request_body.tools[].name - Show call_count column in UI table (sortable) - UI: drop dual_llm option, keep only trusted/blocked * fix: address greptile review feedback - Remove redundant @@index([tool_name]) from schema.prisma (tool_name has @unique which already creates an index) - Replace gen_random_uuid()::text with str(uuid.uuid4()) for portability - Rewrite test_tool_registry_writer.py to mock execute_raw/query_raw (actual implementation) instead of Prisma model methods - Fix test patches in test_tool_management_endpoints.py to target source modules since imports are inside function bodies - Add "Tool Policies" page title to ToolPolicies.tsx * fix: address greptile review round 2 - Replace NOW() with Python datetime parameter in tool_registry_writer (SQLite portability) - Fix cache key collision in tool_policy_guardrail: use null-byte separator instead of colon - Remove type==function filter from request-side tool extraction to match response-side behavior - Clear seen_tool_names on flush so call_count increments per batch cycle not per pod lifetime * fix: address greptile review round 3 - Fix test_seen_names_persist_across_flushes to match actual per-flush-cycle behavior - Update module docstring in tool_discovery_queue.py to accurately describe flush behavior - Add created_at/updated_at to raw SQL INSERT in batch_upsert_tools and update_tool_policy * fix: cache tool policies per tool name not per combination Previously the cache key was built from the full set of tool names in a request, so each unique combination of tools got its own cold cache entry and triggered a separate DB query. With N distinct tools across requests this was effectively a DB hit on every request. Now each tool name is cached individually. Cache hits are checked per tool, only missing tools are fetched from DB in a single batch query, and each result is cached separately. Once a tool's policy is warm, any subsequent request using that tool benefits from the cache regardless of what other tools are in the request. * Update ui/litellm-dashboard/src/components/ToolPolicies.tsx Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/_types.py | 10 + litellm/proxy/db/db_spend_update_writer.py | 151 ++++++- .../tool_discovery_queue.py | 54 +++ litellm/proxy/db/tool_registry_writer.py | 179 ++++++++ .../guardrail_hooks/tool_policy/__init__.py | 16 + .../tool_policy/tool_policy_guardrail.py | 163 +++++++ .../tool_management_endpoints.py | 149 +++++++ litellm/proxy/proxy_server.py | 4 + litellm/proxy/schema.prisma | 20 + litellm/types/tool_management.py | 42 ++ .../test_tool_discovery_queue.py | 75 ++++ .../proxy/db/test_tool_registry_writer.py | 197 +++++++++ .../test_tool_policy_guardrail.py | 181 ++++++++ .../test_tool_management_endpoints.py | 149 +++++++ ui/litellm-dashboard/src/app/page.tsx | 3 + .../src/components/ToolPolicies.tsx | 415 ++++++++++++++++++ .../src/components/leftnav.tsx | 6 + .../src/components/networking.tsx | 54 +++ 19 files changed, 1863 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py create mode 100644 litellm/proxy/db/tool_registry_writer.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py create mode 100644 litellm/proxy/management_endpoints/tool_management_endpoints.py create mode 100644 litellm/types/tool_management.py create mode 100644 tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py create mode 100644 tests/test_litellm/proxy/db/test_tool_registry_writer.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py create mode 100644 ui/litellm-dashboard/src/components/ToolPolicies.tsx diff --git a/litellm/constants.py b/litellm/constants.py index ee79f2fa56f..b1a0021bcc6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -244,6 +244,7 @@ REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) +TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE = int( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f354e28acd7..75b9f91acd9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -183,6 +183,7 @@ class LitellmTableNames(str, enum.Enum): KEY_TABLE_NAME = "LiteLLM_VerificationToken" PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable" MANAGED_FILE_TABLE_NAME = "LiteLLM_ManagedFileTable" + TOOL_TABLE_NAME = "LiteLLM_ToolTable" class Litellm_EntityType(enum.Enum): @@ -4123,6 +4124,15 @@ class SpendUpdateQueueItem(TypedDict, total=False): response_cost: Optional[float] +class ToolDiscoveryQueueItem(TypedDict, total=False): + tool_name: str + origin: Optional[str] # MCP server name or "user_defined" + created_by: Optional[str] + 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 + + class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): unified_file_id: str file_object: Optional[OpenAIFileObject] = None diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d7d6b2b1eb0..edf0cf0d397 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -13,7 +13,17 @@ 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 @@ -23,18 +33,19 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, BaseDailySpendTransaction, - DailyTagSpendTransaction, - DailyOrganizationSpendTransaction, - DailyTeamSpendTransaction, - DailyEndUserSpendTransaction, - DailyUserSpendTransaction, 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, @@ -42,6 +53,9 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( 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: @@ -67,6 +81,7 @@ class DBSpendUpdateWriter: self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) self.pod_lock_manager = PodLockManager() self.spend_update_queue = SpendUpdateQueue() + self.tool_discovery_queue = ToolDiscoveryQueue() self.daily_spend_update_queue = DailySpendUpdateQueue() self.daily_team_spend_update_queue = DailySpendUpdateQueue() self.daily_end_user_spend_update_queue = DailySpendUpdateQueue() @@ -222,6 +237,13 @@ class DBSpendUpdateWriter: ) ) + self._enqueue_tool_registry_upsert( + kwargs=kwargs, + completion_response=completion_response, + hashed_token=hashed_token, + team_id=team_id, + ) + verbose_proxy_logger.debug("Runs spend update on all tables") except Exception: verbose_proxy_logger.error( @@ -237,6 +259,104 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + def _enqueue_tool_registry_upsert( + self, + kwargs: Optional[dict], + completion_response: Optional[Any], + hashed_token: Optional[str] = None, + team_id: Optional[str] = None, + ) -> None: + """ + Extract tool names from the LLM request and response and enqueue them + for upsert into LiteLLM_ToolTable via ToolDiscoveryQueue. + + Handles four sources: + - MCP tools: standard_logging_object.mcp_tool_call_metadata.namespaced_tool_name + - Response tool_calls (OpenAI / Anthropic pass-through converted to OpenAI format): + completion_response.choices[].message.tool_calls[].function.name + - Request tools array (OpenAI format): kwargs["tools"][].function.name + - Request tools array (Anthropic /messages format): kwargs["passthrough_logging_payload"] + ["request_body"]["tools"][].name + """ + try: + if kwargs is None: + return + + # Extract key_alias from kwargs metadata if available + key_alias: Optional[str] = None + _litellm_params = kwargs.get("litellm_params") or {} + _metadata = _litellm_params.get("metadata") or {} + key_alias = _metadata.get("user_api_key_alias") or None + + def _enqueue(tool_name: str, origin: str = "user_defined") -> None: + self.tool_discovery_queue.add_update( + ToolDiscoveryQueueItem( + tool_name=tool_name, + origin=origin, + key_hash=hashed_token, + team_id=team_id, + key_alias=key_alias, + ) + ) + + # --- 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") + if mcp_metadata and isinstance(mcp_metadata, dict): + 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") + + # --- Tools from request body (OpenAI format: tools[].function.name) --- + request_tools = kwargs.get("tools") or [] + for tool_def in request_tools: + if not isinstance(tool_def, dict): + continue + fn = tool_def.get("function") or {} + name = fn.get("name") if isinstance(fn, dict) else None + if name: + _enqueue(name) + + # --- Tools from Anthropic /messages pass-through request body + # (Anthropic format: tools[].name, no "function" wrapper) --- + passthrough_payload = kwargs.get("passthrough_logging_payload") or {} + request_body = ( + passthrough_payload.get("request_body") + if isinstance(passthrough_payload, dict) + else None + ) or {} + for tool_def in request_body.get("tools") or []: + if not isinstance(tool_def, dict): + continue + name = tool_def.get("name") + if name: + _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"): + for choice in completion_response.choices or []: + message = getattr(choice, "message", None) + if message is None: + continue + tool_calls = getattr(message, "tool_calls", None) + if not tool_calls: + continue + for tc in tool_calls: + fn = getattr(tc, "function", None) + if fn is None: + continue + tool_name = getattr(fn, "name", None) + if tool_name: + _enqueue(tool_name) + except Exception as e: + verbose_proxy_logger.debug( + "_enqueue_tool_registry_upsert error (non-blocking): %s", e + ) + async def _update_key_db( self, response_cost: Optional[float], @@ -752,6 +872,25 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) + ################## Tool Registry Upserts ################## + await self._flush_tool_discovery_queue(prisma_client=prisma_client) + + async def _flush_tool_discovery_queue( + self, + prisma_client: PrismaClient, + ) -> None: + """Flush ToolDiscoveryQueue and batch-upsert new tools into LiteLLM_ToolTable.""" + from litellm.proxy.db.tool_registry_writer import batch_upsert_tools + + try: + items = self.tool_discovery_queue.flush() + if items: + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: + verbose_proxy_logger.debug( + "_flush_tool_discovery_queue error (non-blocking): %s", e + ) + async def _commit_spend_updates_to_db( # noqa: PLR0915 self, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py new file mode 100644 index 00000000000..16a3ada40f2 --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -0,0 +1,54 @@ +""" +In-memory buffer for tool registry upserts. + +Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue +uses set-deduplication: each unique tool_name is only queued once per flush +cycle (~30s). The seen-set is cleared on every flush so that call_count +increments in subsequent cycles rather than stopping after the first flush. +""" + +from typing import List, Set + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ToolDiscoveryQueueItem + + +class ToolDiscoveryQueue: + """ + In-memory buffer for tool registry upserts. + + Deduplicates by tool_name within each flush cycle: a tool is only queued + once per ~30s batch, so call_count increments once per flush cycle the + tool appears in (not once per invocation, but not once per pod lifetime + either). The seen-set is cleared on flush so subsequent batches can + re-count the same tool. + """ + + def __init__(self) -> None: + self._seen_tool_names: Set[str] = set() + self._pending: List[ToolDiscoveryQueueItem] = [] + + def add_update(self, item: ToolDiscoveryQueueItem) -> None: + """Enqueue a tool discovery item if tool_name has not been seen before.""" + tool_name = item.get("tool_name", "") + if not tool_name: + return + if tool_name in self._seen_tool_names: + verbose_proxy_logger.debug( + "ToolDiscoveryQueue: skipping already-seen tool %s", tool_name + ) + return + self._seen_tool_names.add(tool_name) + self._pending.append(item) + verbose_proxy_logger.debug( + "ToolDiscoveryQueue: queued new tool %s (origin=%s)", + tool_name, + item.get("origin"), + ) + + def flush(self) -> List[ToolDiscoveryQueueItem]: + """Return and clear all pending items. Resets seen-set so the next + flush cycle can re-count the same tools.""" + items, self._pending = self._pending, [] + self._seen_tool_names.clear() + return items diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py new file mode 100644 index 00000000000..4e0a8095a08 --- /dev/null +++ b/litellm/proxy/db/tool_registry_writer.py @@ -0,0 +1,179 @@ +""" +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. +""" + +import uuid +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Dict, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ToolDiscoveryQueueItem +from litellm.types.tool_management import LiteLLM_ToolTableRow, ToolCallPolicy + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +def _row_to_model(row: dict) -> LiteLLM_ToolTableRow: + 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"), + 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"), + created_at=row.get("created_at"), + updated_at=row.get("updated_at"), + created_by=row.get("created_by"), + updated_by=row.get("updated_by"), + ) + + +async def batch_upsert_tools( + prisma_client: "PrismaClient", + items: List[ToolDiscoveryQueueItem], +) -> None: + """ + Batch-upsert tool registry rows via raw SQL. + + On first insert: sets call_policy = "untrusted" (schema default), call_count = 1. + On conflict: increments call_count; preserves existing call_policy. + """ + if not items: + return + try: + data = [item for item in items if item.get("tool_name")] + if not data: + return + for item in data: + tool_name = item.get("tool_name", "") + origin = item.get("origin") or "user_defined" + created_by = item.get("created_by") or "system" + 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, + ) + 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) + + +async def list_tools( + prisma_client: "PrismaClient", + call_policy: Optional[ToolCallPolicy] = None, +) -> List[LiteLLM_ToolTableRow]: + """Return all tools, optionally filtered by call_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', + ) + 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) + return [] + + +async def get_tool( + prisma_client: "PrismaClient", + tool_name: str, +) -> 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, + ) + if not rows: + return None + return _row_to_model(rows[0]) + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer get_tool error: %s", e) + return None + + +async def update_tool_policy( + prisma_client: "PrismaClient", + tool_name: str, + call_policy: ToolCallPolicy, + updated_by: Optional[str], +) -> Optional[LiteLLM_ToolTableRow]: + """Update the call_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, + ) + 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) + return None + + +async def get_tools_by_names( + prisma_client: "PrismaClient", + tool_names: List[str], +) -> Dict[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. + """ + 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, + ) + return {row["tool_name"]: row["call_policy"] for row in rows} + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer get_tools_by_names error: %s", e) + return {} diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py new file mode 100644 index 00000000000..5a43006e23c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py @@ -0,0 +1,16 @@ +import litellm +from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail): + from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( + ToolPolicyGuardrail, + ) + + _callback = ToolPolicyGuardrail( + guardrail_name=guardrail.get("guardrail_name", "tool_policy"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback 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 new file mode 100644 index 00000000000..87558566c42 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -0,0 +1,163 @@ +""" +Tool Policy Guardrail + +Reads call_policy from LiteLLM_ToolTable and enforces it on LLM requests/responses. + +Policy values: + "trusted" - allow through (no action) + "untrusted" - allow through (no action; default for newly discovered tools) + "blocked" - raise HTTPException, preventing the tool call + "dual_llm" - (Phase 3) send to second LLM for verification; currently treated as allowed + +Configuration in proxy config YAML: + guardrails: + - guardrail_name: "tool_policy" + 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 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.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +GUARDRAIL_NAME = "tool_policy" + + +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. + """ + + def __init__(self, **kwargs: Any) -> None: + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.during_call, + ] + super().__init__(**kwargs) + self._policy_cache: DualCache = DualCache() + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + 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". + """ + if input_type == "request": + tools = inputs.get("tools") or [] + tool_names = [ + t["function"]["name"] + for t in tools + if isinstance(t, dict) + and isinstance(t.get("function"), dict) + and t["function"].get("name") + ] + else: # response + tool_calls = inputs.get("tool_calls") or [] + tool_names = [] + for tc in tool_calls: + fn = None + if isinstance(tc, dict): + fn = (tc.get("function") or {}).get("name") + elif hasattr(tc, "function"): + fn = getattr(tc.function, "name", None) + if fn: + tool_names.append(fn) + + if not tool_names: + return inputs + + policy_map = await self._get_policies_cached(tool_names) + + 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 + ) + raise HTTPException( + status_code=400, + detail={ + "error": "Violated tool policy", + "blocked_tools": blocked, + "message": f"Tool(s) {blocked} are blocked by policy.", + }, + ) + + 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/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py new file mode 100644 index 00000000000..89880c9a4ec --- /dev/null +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -0,0 +1,149 @@ +""" +TOOL POLICY MANAGEMENT + +All /tool management endpoints + +GET /v1/tool/list - List all discovered tools and their policies +GET /v1/tool/{tool_name} - Get a single tool's details +POST /v1/tool/policy - Update the call_policy for a tool +""" + +from typing import Optional + +from fastapi import APIRouter, Depends, HTTPException + +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, + ToolListResponse, + ToolPolicyUpdateRequest, + ToolPolicyUpdateResponse, +) + +router = APIRouter() + + +@router.get( + "/v1/tool/list", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolListResponse, +) +async def list_tools( + call_policy: Optional[ToolCallPolicy] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + List all auto-discovered tools and their call policies. + + Parameters: + - call_policy: Optional filter — one of "trusted", "untrusted", "dual_llm", "blocked" + """ + from litellm.proxy.db.tool_registry_writer import list_tools as db_list_tools + 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: + tools = await db_list_tools(prisma_client=prisma_client, call_policy=call_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}", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_ToolTableRow, +) +async def get_tool( + tool_name: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + 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 + + 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" + ) + return tool + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error getting tool: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/v1/tool/policy", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolPolicyUpdateResponse, +) +async def update_tool_policy( + data: ToolPolicyUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Set the call policy for a tool. + + 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. + """ + from litellm.proxy.db.tool_registry_writer import ( + update_tool_policy as db_update_tool_policy, + ) + 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: + 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, + ) + if updated is None: + raise HTTPException( + status_code=500, detail=f"Failed to update policy for tool '{data.tool_name}'" + ) + return ToolPolicyUpdateResponse( + tool_name=updated.tool_name, + call_policy=updated.call_policy, + updated=True, + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error updating tool policy: %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 905a4be35e4..36b6bcc7707 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -410,6 +410,9 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team, validate_membership, ) +from litellm.proxy.management_endpoints.tool_management_endpoints import ( + router as tool_management_router, +) from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) @@ -12882,6 +12885,7 @@ app.include_router(budget_management_router) app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) +app.include_router(tool_management_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 50c0a55a875..23917cf7c7f 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1051,6 +1051,26 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } +// Global tool registry - auto-discovered from LLM responses; admins set call_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? + + @@index([call_policy]) + @@index([team_id]) +} + //Unified Access Groups table for storing unified access groups model LiteLLM_AccessGroupTable { access_group_id String @id @default(uuid()) diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py new file mode 100644 index 00000000000..8704ff27759 --- /dev/null +++ b/litellm/types/tool_management.py @@ -0,0 +1,42 @@ +""" +Pydantic models for Tool Policy management endpoints. +""" + +from datetime import datetime +from typing import Dict, List, Literal, Optional + +from pydantic import BaseModel + +ToolCallPolicy = Literal["trusted", "untrusted", "dual_llm", "blocked"] + + +class LiteLLM_ToolTableRow(BaseModel): + tool_id: str + tool_name: str + origin: Optional[str] = None + call_policy: ToolCallPolicy = "untrusted" + call_count: int = 0 + assignments: Optional[Dict] = None + key_hash: Optional[str] = None + team_id: Optional[str] = None + key_alias: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_by: Optional[str] = None + + +class ToolListResponse(BaseModel): + tools: List[LiteLLM_ToolTableRow] + total: int + + +class ToolPolicyUpdateRequest(BaseModel): + tool_name: str + call_policy: ToolCallPolicy + + +class ToolPolicyUpdateResponse(BaseModel): + tool_name: str + call_policy: ToolCallPolicy + updated: bool diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py new file mode 100644 index 00000000000..defdb3834d8 --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -0,0 +1,75 @@ +""" +Unit tests for ToolDiscoveryQueue. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, +) + + +@pytest.fixture +def queue(): + return ToolDiscoveryQueue() + + +def test_add_single_tool(queue): + queue.add_update({"tool_name": "my_tool", "origin": "user_defined"}) + items = queue.flush() + assert len(items) == 1 + assert items[0]["tool_name"] == "my_tool" + assert items[0]["origin"] == "user_defined" + + +def test_deduplication_same_name(queue): + """Adding the same tool_name twice should only keep the first.""" + queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"}) + queue.add_update({"tool_name": "tool_a", "origin": "user_defined"}) + items = queue.flush() + assert len(items) == 1 + assert items[0]["origin"] == "mcp_server" # first wins + + +def test_deduplication_different_names(queue): + queue.add_update({"tool_name": "tool_a"}) + queue.add_update({"tool_name": "tool_b"}) + items = queue.flush() + assert len(items) == 2 + names = {i["tool_name"] for i in items} + assert names == {"tool_a", "tool_b"} + + +def test_flush_clears_pending(queue): + queue.add_update({"tool_name": "tool_x"}) + items1 = queue.flush() + assert len(items1) == 1 + items2 = queue.flush() + assert len(items2) == 0 + + +def test_seen_names_reset_after_flush(queue): + """Seen-set is cleared on flush so the same tool can re-enter the next cycle.""" + queue.add_update({"tool_name": "tool_a"}) + queue.flush() + queue.add_update({"tool_name": "tool_a"}) # same tool, new cycle + items = queue.flush() + assert len(items) == 1 + assert items[0]["tool_name"] == "tool_a" + + +def test_empty_tool_name_ignored(queue): + queue.add_update({"tool_name": ""}) + queue.add_update({"tool_name": None}) # type: ignore[arg-type] + items = queue.flush() + assert len(items) == 0 + + +def test_flush_returns_list(queue): + result = queue.flush() + assert isinstance(result, list) diff --git a/tests/test_litellm/proxy/db/test_tool_registry_writer.py b/tests/test_litellm/proxy/db/test_tool_registry_writer.py new file mode 100644 index 00000000000..44f9e32058a --- /dev/null +++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py @@ -0,0 +1,197 @@ +""" +Unit tests for tool_registry_writer.py — uses a mock prisma client +that exposes execute_raw / query_raw (matching the actual raw-SQL implementation). +""" + +import os +import sys +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.tool_registry_writer import ( + batch_upsert_tools, + get_tool, + 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 = { + "tool_id": "uuid-1", + "tool_name": "my_tool", + "origin": "user_defined", + "call_policy": "untrusted", + "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": None, + } + rows = query_rows if query_rows is not None else [default_row] + + prisma = MagicMock() + prisma.db.execute_raw = AsyncMock(return_value=None) + prisma.db.query_raw = AsyncMock(return_value=rows) + return prisma + + +@pytest.mark.asyncio +async def test_batch_upsert_tools_calls_execute_raw(): + 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 + + +@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() + + +@pytest.mark.asyncio +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() + + +@pytest.mark.asyncio +async def test_batch_upsert_multiple_tools_calls_execute_raw_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 + + +@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]) + 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() + + +@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 + + +@pytest.mark.asyncio +async def test_get_tool_found(): + prisma = _make_prisma() + 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() + + +@pytest.mark.asyncio +async def test_get_tool_not_found(): + prisma = _make_prisma(query_rows=[]) + 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") + 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 + + +@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"}, + ] + prisma = _make_prisma(query_rows=rows) + result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) + assert result == {"tool_a": "trusted", "tool_b": "blocked"} + + +@pytest.mark.asyncio +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() 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 new file mode 100644 index 00000000000..c6a81efbf0b --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py @@ -0,0 +1,181 @@ +""" +Unit tests for ToolPolicyGuardrail. +""" + +import os +import sys +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +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.types.guardrails import GuardrailEventHooks + + +@pytest.fixture +def guardrail(): + return ToolPolicyGuardrail() + + +# --- helpers --- + +def _tool_request_inputs(tool_names: list) -> dict: + return { + "tools": [ + {"type": "function", "function": {"name": name, "description": ""}} + for name in tool_names + ] + } + + +def _tool_response_inputs(tool_names: list) -> dict: + return { + "tool_calls": [ + {"type": "function", "function": {"name": name}} + for name in tool_names + ] + } + + +# --- tests --- + + +def test_guardrail_supports_pre_and_post_call(guardrail): + hooks = guardrail.supported_event_hooks + assert GuardrailEventHooks.pre_call in hooks + assert GuardrailEventHooks.post_call in hooks + + +@pytest.mark.asyncio +async def test_no_tools_in_request_passes_through(guardrail): + inputs: Any = {"tools": []} + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_no_tool_calls_in_response_passes_through(guardrail): + inputs: Any = {"tool_calls": []} + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + +@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)): + inputs: Any = _tool_request_inputs(["search", "read_file"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@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)): + inputs: Any = _tool_request_inputs(["dangerous_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert exc_info.value.status_code == 400 + assert "dangerous_tool" in exc_info.value.detail["blocked_tools"] + + +@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)): + inputs: Any = _tool_response_inputs(["exfil_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert exc_info.value.status_code == 400 + assert "exfil_tool" in exc_info.value.detail["blocked_tools"] + + +@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)): + inputs: Any = _tool_request_inputs(["safe_tool", "bad_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + blocked = exc_info.value.detail["blocked_tools"] + assert "bad_tool" in blocked + assert "safe_tool" not in blocked + + +@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={})): + inputs: Any = _tool_request_inputs(["unknown_tool"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@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"} + 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(), + ): + # 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 == {} + + +@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)): + fn = MagicMock() + fn.name = "obj_tool" + tc = MagicMock() + tc.function = fn + inputs: Any = {"tool_calls": [tc]} + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py new file mode 100644 index 00000000000..6f1d373fdee --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -0,0 +1,149 @@ +""" +Unit tests for tool management endpoints (/v1/tool/*). +Uses FastAPI TestClient with mocked DB functions. + +Patches target the source modules (litellm.proxy.db.tool_registry_writer.* +and litellm.proxy.proxy_server.prisma_client) because the endpoint code +imports these inside function bodies to avoid circular imports. +""" + +import os +import sys +from datetime import datetime, timezone +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.management_endpoints.tool_management_endpoints import router +from litellm.types.tool_management import LiteLLM_ToolTableRow + +# --- helpers --- + + +def _make_tool_row( + tool_name: str = "my_tool", + call_policy: str = "untrusted", + origin: Optional[str] = None, +) -> LiteLLM_ToolTableRow: + now = datetime.now(timezone.utc) + return LiteLLM_ToolTableRow( + tool_id="uuid-1", + tool_name=tool_name, + origin=origin, + call_policy=call_policy, # type: ignore[arg-type] + assignments={}, + created_at=now, + updated_at=now, + ) + + +def _make_app() -> FastAPI: + """Build a minimal FastAPI app with the tool management router.""" + app = FastAPI() + app.include_router(router) + return app + + +# Stub the auth dependency so we don't need a real proxy running. +def _override_auth(): + from litellm.proxy._types import UserAPIKeyAuth + + return UserAPIKeyAuth(api_key="sk-test", user_id="admin") + + +# A real (non-None) prisma stub for truthiness checks. +_MOCK_PRISMA = MagicMock() + + +# --- test class --- + + +class TestToolManagementEndpoints: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch( + "litellm.proxy.db.tool_registry_writer.list_tools", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_list_tools_returns_200(self, mock_db_list): + mock_db_list.return_value = [_make_tool_row()] + + resp = self.client.get("/v1/tool/list") + assert resp.status_code == 200 + body = resp.json() + assert body["total"] == 1 + assert body["tools"][0]["tool_name"] == "my_tool" + + @patch( + "litellm.proxy.db.tool_registry_writer.list_tools", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_list_tools_with_policy_filter(self, mock_db_list): + mock_db_list.return_value = [_make_tool_row(call_policy="blocked")] + + resp = self.client.get("/v1/tool/list?call_policy=blocked") + assert resp.status_code == 200 + assert resp.json()["tools"][0]["call_policy"] == "blocked" + + @patch( + "litellm.proxy.db.tool_registry_writer.get_tool", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_get_tool_found(self, mock_db_get): + mock_db_get.return_value = _make_tool_row(tool_name="tool_a") + + resp = self.client.get("/v1/tool/tool_a") + assert resp.status_code == 200 + assert resp.json()["tool_name"] == "tool_a" + + @patch( + "litellm.proxy.db.tool_registry_writer.get_tool", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_get_tool_not_found_returns_404(self, mock_db_get): + mock_db_get.return_value = None + + resp = self.client.get("/v1/tool/nonexistent", follow_redirects=True) + assert resp.status_code == 404 + + @patch( + "litellm.proxy.db.tool_registry_writer.update_tool_policy", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_update_tool_policy_blocked(self, mock_db_update): + mock_db_update.return_value = _make_tool_row(call_policy="blocked") + + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "call_policy": "blocked"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["call_policy"] == "blocked" + assert body["updated"] is True + + @patch("litellm.proxy.proxy_server.prisma_client", None) + def test_list_tools_no_db_returns_500(self): + resp = self.client.get("/v1/tool/list") + assert resp.status_code == 500 + + def test_update_tool_policy_invalid_policy_returns_422(self): + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "call_policy": "invalid_value"}, + ) + assert resp.status_code == 422 diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index fb749d7afb0..258c2ccb0e0 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -38,6 +38,7 @@ import Usage from "@/components/usage"; import UserDashboard from "@/components/user_dashboard"; import { AccessGroupsPage } from "@/components/AccessGroups/AccessGroupsPage"; import VectorStoreManagement from "@/components/vector_store_management"; +import ToolPolicies from "@/components/ToolPolicies"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; @@ -548,6 +549,8 @@ function CreateKeyPageContent() { ) : page == "vector-stores" ? ( + ) : page == "tool-policies" ? ( + ) : page == "guardrails-monitor" ? ( ) : page == "new_usage" ? ( diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx new file mode 100644 index 00000000000..0e3f5434e7f --- /dev/null +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -0,0 +1,415 @@ +"use client"; + +import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; +import { Select, Switch, Tooltip } from "antd"; +import { Select, Tooltip } from "antd"; +import { + Table, + TableHead, + TableHeaderCell, + TableBody, + TableRow, + TableCell, +} from "@tremor/react"; +import { TimeCell } from "./view_logs/time_cell"; +import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; +import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; +import FilterComponent, { FilterOption } from "./molecules/filter"; +import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking"; + +const POLICY_OPTIONS = [ + { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, + { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, +] as const; + +type PolicyValue = "trusted" | "blocked"; + +const policyStyle = (p: string) => + POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; + +type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at" | "call_count"; + +interface FilterValues { + [key: string]: string; +} + +interface ToolPoliciesProps { + accessToken: string | null; + userRole?: string; +} + +const PolicySelect: React.FC<{ + value: string; + toolName: string; + saving: boolean; + onChange: (toolName: string, policy: string) => void; +}> = ({ value, toolName, saving, onChange }) => { + const style = policyStyle(value); + return ( + { setSearchTerm(e.target.value); setCurrentPage(1); }} + /> + + + + + +
+ Live Tail + +
+ + + + +
+ + Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results + + Page {currentPage} of {totalPages} +
+ + +
+
+ + + {/* Filter row */} +
+ +
+ + + {/* Auto-refresh banner */} + {isLiveTail && ( +
+ Auto-refreshing every 15 seconds + +
+ )} + + {error && ( +
{error}
+ )} + + {/* Table */} + + + + + + + + + Key Hash + + Origin + + + + {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.team_id ?? "-"} + + + + + + {tool.key_hash ?? "-"} + + + + + + {tool.key_alias ?? "-"} + + + + + {tool.origin ?? "-"} + + + + )) + )} + +
+ + {/* Bottom pagination (only when > 1 page) */} + {totalPages > 1 && ( +
+ Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length} +
+ + +
+
+ )} + + + ); +}; + +export default ToolPolicies; diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index da3ca2a8bae..2cbeb22ec81 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -134,6 +134,12 @@ const menuGroups: MenuGroup[] = [ label: "Vector Stores", icon: , }, + { + key: "tool-policies", + page: "tool-policies", + label: "Tool Policies", + icon: , + }, ], }, ], diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 6ffd744cce9..8536a584ee1 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -9854,3 +9854,57 @@ export const checkGdprCompliance = async ( } return response.json(); }; + +export interface ToolRow { + tool_id: string; + tool_name: string; + origin?: string; + call_policy: string; + call_count?: number; + assignments?: Record; + key_hash?: string; + team_id?: string; + key_alias?: string; + created_at?: string; + updated_at?: string; + created_by?: string; + updated_by?: string; +} + +export const fetchToolsList = async (accessToken: string): Promise => { + const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/list` : `/v1/tool/list`; + const response = await fetch(url, { + method: "GET", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + if (!response.ok) { + const errorData = await response.text(); + throw new Error(errorData); + } + const data = await response.json(); + return data.tools ?? []; +}; + +export const updateToolPolicy = async ( + accessToken: string, + toolName: string, + callPolicy: string +): Promise => { + const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/policy` : `/v1/tool/policy`; + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ tool_name: toolName, call_policy: callPolicy }), + }); + if (!response.ok) { + const errorData = await response.text(); + throw new Error(errorData); + } + return response.json(); +}; From 360643e21315eb02ca790eb22effb87a6d48167b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 16:40:04 -0800 Subject: [PATCH 039/132] [Feat] UI - Allow using AI to understand Usage patterns (#22042) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add Ask AI chat component to Usage page - Create UsageAIChatModal component with streaming chat interface - Integrate with existing model hub for model selection - Pass usage data context (spend, models, providers, keys) to AI - Add Ask AI button next to Export Data button in global view - Add tests for the new component and integration Co-authored-by: Ishaan Jaff * Convert Ask AI from modal to right-side sliding panel - Replace UsageAIChatModal with UsageAIChatPanel - Panel slides in from right side, usage page stays visible - Full-height panel with header, model selector, chat area, and input - Smooth CSS transition for open/close animation - Update tests for new panel component (34 tests passing) Co-authored-by: Ishaan Jaff * Remove build output directory from tracking Co-authored-by: Ishaan Jaff * Add backend AI usage chat endpoint with tool calling Backend: - New /usage/ai/chat SSE streaming endpoint - AI agent has get_usage_data tool that queries /user/daily/activity/aggregated - Follows same architecture as policy AI suggest (litellm.acompletion + tools) - Non-admin users are restricted to their own data - 12 backend unit tests Frontend: - Panel now calls /usage/ai/chat backend endpoint via SSE - Removed direct OpenAI client calls from frontend - Added usageAiChatStream networking function following enrichPolicyTemplateStream pattern Co-authored-by: Ishaan Jaff * Make model selection optional, default to gpt-4o-mini on backend Co-authored-by: Ishaan Jaff * Add team/tag tools, status indicators, and improved AI agent - AI agent now has 3 tools: get_usage_data, get_team_usage_data, get_tag_usage_data - Stream status events (Thinking... Fetching... Analyzing...) to UI - Frontend shows spinner + status text during tool execution - Better system prompt guiding tool selection - Entity summariser for team/tag data with ranked breakdowns - 13 backend tests, 34 frontend tests passing Co-authored-by: Ishaan Jaff * Fix: inject today's date into system prompt so AI resolves relative dates correctly Co-authored-by: Ishaan Jaff * Show tool calls as distinct steps + render markdown in responses - Backend emits tool_call events with tool_name, label, args, and status - Frontend shows each tool call as a step with ✓/spinner/✗ indicator - Tool call steps show icon, label, date range, and filters - AI responses rendered with ReactMarkdown (bold, lists, tables, code) - Cursor-like UX: Thinking → tool calls → Analyzing → streamed answer Co-authored-by: Ishaan Jaff * Refactor backend for code quality: proper types, constants, all functions ≤50 LOC - TypedDict for SSE events (SSEStatusEvent, SSEToolCallEvent, etc.) and ToolHandler - Constants for table names, entity fields, temperature, page sizes, top-N limits - Shared _query_activity() eliminates duplicated fetch logic - _accumulate_breakdown() + _ranked_lines() replace inline aggregation loops - Extracted _process_tool_call() and _stream_final_response() from main stream fn - Black + Ruff clean, all 15 functions verified ≤50 LOC - Replaced Tremor Button with Antd Button in panel (Tremor deprecated per AGENTS.md) Co-authored-by: Ishaan Jaff * Address greptile review: security fixes and input validation - Restrict team/tag tools to admin-only users (non-admins only get get_usage_data) - Constrain ChatMessage.role to Literal['user', 'assistant'] to prevent system prompt injection - Add test for base tools restriction (non-admin gets 1 tool, admin gets 3) - Issues 3 (unused imports) and 4 (inline datetime) were already fixed in prior commit Co-authored-by: Ishaan Jaff * Address greptile round 2: sanitize errors, defense-in-depth allowlist, revert tsconfig - Sanitize error messages: generic 'An internal error occurred' sent to client, full exception logged server-side via verbose_proxy_logger - Defense-in-depth: _process_tool_call validates fn_name against role-based allowlist before dispatch (even though LLM only receives allowed tools) - Revert tsconfig.json jsx back to 'preserve' (Next.js recommended default) Co-authored-by: Ishaan Jaff * Role-scoped system prompt + additional test coverage - System prompt is now role-aware: admin sees all 3 tool descriptions, non-admin only sees get_usage_data (consistent with tool filtering) - Added tests: non-admin prompt excludes team/tag tools, date injection - 15 backend tests, 34 frontend tests passing Co-authored-by: Ishaan Jaff * Fix LLM arg validation + cap conversation size at 20 messages - _resolve_fetch_kwargs uses .get() with ValueError for missing dates (handles malformed LLM tool arguments gracefully) - MAX_CHAT_MESSAGES = 20 constant; backend truncates to last 20 - Frontend also sends only last 20 messages per request - Prevents excessive token usage and context-length errors Co-authored-by: Ishaan Jaff --------- Co-authored-by: Cursor Agent Co-authored-by: Ishaan Jaff --- .../usage_endpoints/__init__.py | 9 + .../usage_endpoints/ai_usage_chat.py | 578 ++++++++++++++++++ .../usage_endpoints/endpoints.py | 65 ++ litellm/proxy/proxy_server.py | 2 + .../usage_endpoints/__init__.py | 0 .../usage_endpoints/test_ai_usage_chat.py | 402 ++++++++++++ ui/litellm-dashboard/package-lock.json | 151 +---- .../components/UsageAIChatPanel.test.tsx | 85 +++ .../UsagePage/components/UsageAIChatPanel.tsx | 402 ++++++++++++ .../components/UsagePageView.test.tsx | 26 + .../UsagePage/components/UsagePageView.tsx | 51 +- .../src/components/networking.tsx | 76 +++ 12 files changed, 1704 insertions(+), 143 deletions(-) create mode 100644 litellm/proxy/management_endpoints/usage_endpoints/__init__.py create mode 100644 litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py create mode 100644 litellm/proxy/management_endpoints/usage_endpoints/endpoints.py create mode 100644 tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py create mode 100644 ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx create mode 100644 ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx diff --git a/litellm/proxy/management_endpoints/usage_endpoints/__init__.py b/litellm/proxy/management_endpoints/usage_endpoints/__init__.py new file mode 100644 index 00000000000..6e68dcd2a2e --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/__init__.py @@ -0,0 +1,9 @@ +""" +Usage endpoints package. + +Re-exports the router from endpoints module. +""" + +from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( # noqa: F401 + router, +) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py new file mode 100644 index 00000000000..f156be7d2cc --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -0,0 +1,578 @@ +""" +AI Usage Chat - uses LLM tool calling to answer questions about +usage/spend data by querying the aggregated daily activity endpoints. +""" + +import json +from datetime import date +from typing import Any, AsyncIterator, Callable, Dict, List, Literal, Optional + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) + +from typing_extensions import TypedDict + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +USAGE_AI_TEMPERATURE = 0.2 + +TABLE_DAILY_USER_SPEND = "litellm_dailyuserspend" +TABLE_DAILY_TEAM_SPEND = "litellm_dailyteamspend" +TABLE_DAILY_TAG_SPEND = "litellm_dailytagspend" + +ENTITY_FIELD_USER = "user_id" +ENTITY_FIELD_TEAM = "team_id" +ENTITY_FIELD_TAG = "tag" + +PAGINATED_PAGE_SIZE = 200 +MAX_CHAT_MESSAGES = 20 +TOP_N_MODELS = 15 +TOP_N_PROVIDERS = 10 +TOP_N_KEYS = 10 + +# --------------------------------------------------------------------------- +# Types +# --------------------------------------------------------------------------- + + +class SSEStatusEvent(TypedDict): + type: Literal["status"] + message: str + + +class SSEToolCallEvent(TypedDict, total=False): + type: Literal["tool_call"] + tool_name: str + tool_label: str + arguments: Dict[str, str] + status: Literal["running", "complete", "error"] + error: str + + +class SSEChunkEvent(TypedDict): + type: Literal["chunk"] + content: str + + +class SSEDoneEvent(TypedDict): + type: Literal["done"] + + +class SSEErrorEvent(TypedDict): + type: Literal["error"] + message: str + + +SSEEvent = ( + SSEStatusEvent | SSEToolCallEvent | SSEChunkEvent | SSEDoneEvent | SSEErrorEvent +) + + +class ToolHandler(TypedDict): + fetch: Callable[..., Any] + summarise: Callable[[Dict[str, Any]], str] + label: str + + +# --------------------------------------------------------------------------- +# Tool definitions (OpenAI function-calling schema) +# --------------------------------------------------------------------------- + +_DATE_PARAMS = { + "start_date": {"type": "string", "description": "Start date in YYYY-MM-DD format"}, + "end_date": {"type": "string", "description": "End date in YYYY-MM-DD format"}, +} + +_TOOL_USAGE = { + "type": "function", + "function": { + "name": "get_usage_data", + "description": ( + "Fetch aggregated global usage/spend data. Returns daily spend, " + "token counts, request counts, and breakdowns by model, provider, " + "and API key. Use for overall spend, top models, top providers." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "user_id": { + "type": "string", + "description": "Optional user ID filter. Omit for global view.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +_TOOL_TEAM = { + "type": "function", + "function": { + "name": "get_team_usage_data", + "description": ( + "Fetch usage/spend data broken down by team. Use for questions " + "like 'which team spends the most' or 'show me team X usage'." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "team_ids": { + "type": "string", + "description": "Optional comma-separated team IDs. Omit for all teams.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +_TOOL_TAG = { + "type": "function", + "function": { + "name": "get_tag_usage_data", + "description": ( + "Fetch usage/spend data broken down by tag. Tags are labels " + "attached to requests (features, environments, credentials)." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "tags": { + "type": "string", + "description": "Optional comma-separated tag names. Omit for all tags.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +TOOLS_BASE = [_TOOL_USAGE] +TOOLS_ADMIN = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG] + + +def get_tools_for_role(is_admin: bool) -> List[Dict[str, Any]]: + """Return the tool list appropriate for the user's role.""" + return TOOLS_ADMIN if is_admin else TOOLS_BASE + + +_SYSTEM_PROMPT_BASE = ( + "You are an AI assistant embedded in the LiteLLM Usage dashboard. " + "You help users understand their LLM API spend and usage data.\n\n" + "ALWAYS call the appropriate tool(s) first to fetch data before answering. " + "You may call multiple tools if the question spans different dimensions.\n\n" + "Guidelines:\n" + "- Be concise and specific. Use exact numbers from the data.\n" + "- Format costs as dollar amounts (e.g. $12.34).\n" + "- When comparing entities, show a ranked list.\n" + "- If data is empty or no results found, say so clearly.\n" + "- Do not hallucinate data — only use what the tools return.\n" + "- Today's date will be provided below. Use it to interpret relative dates " + "like 'this week', 'this month', 'last 7 days', etc." +) + +_TOOL_DESCRIPTIONS_ADMIN = ( + "You have access to these tools:\n" + "- `get_usage_data`: Global/user-level usage (spend, models, providers, API keys)\n" + "- `get_team_usage_data`: Team-level usage breakdown\n" + "- `get_tag_usage_data`: Tag-level usage breakdown\n\n" +) + +_TOOL_DESCRIPTIONS_BASE = ( + "You have access to this tool:\n" + "- `get_usage_data`: Your usage data (spend, models, providers, API keys)\n\n" +) + + +def _build_system_prompt(is_admin: bool) -> str: + """Build role-appropriate system prompt with today's date.""" + tool_desc = _TOOL_DESCRIPTIONS_ADMIN if is_admin else _TOOL_DESCRIPTIONS_BASE + return ( + f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}" + f"Today's date: {date.today().isoformat()}" + ) + + +# keep a public reference for test assertions +SYSTEM_PROMPT = _SYSTEM_PROMPT_BASE + +# --------------------------------------------------------------------------- +# Data fetchers +# --------------------------------------------------------------------------- + + +def _parse_csv_ids(raw: Optional[str]) -> Optional[List[str]]: + if not raw: + return None + return [t.strip() for t in raw.split(",") if t.strip()] + + +async def _query_activity( + table_name: str, + entity_id_field: str, + entity_id: Optional[Any], + start_date: str, + end_date: str, + *, + use_aggregated: bool = False, +) -> SpendAnalyticsPaginatedResponse: + """Shared helper that calls the daily activity query layer.""" + from litellm.proxy.management_endpoints.common_daily_activity import ( + get_daily_activity, + get_daily_activity_aggregated, + ) + from litellm.proxy.proxy_server import prisma_client + + if use_aggregated: + return await get_daily_activity_aggregated( + prisma_client=prisma_client, + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + entity_metadata_field=None, + start_date=start_date, + end_date=end_date, + model=None, + api_key=None, + ) + return await get_daily_activity( + prisma_client=prisma_client, + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + entity_metadata_field=None, + start_date=start_date, + end_date=end_date, + model=None, + api_key=None, + page=1, + page_size=PAGINATED_PAGE_SIZE, + ) + + +async def _fetch_usage_data( + start_date: str, end_date: str, user_id: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_USER_SPEND, + ENTITY_FIELD_USER, + user_id, + start_date, + end_date, + use_aggregated=True, + ) + return resp.model_dump(mode="json") + + +async def _fetch_team_usage_data( + start_date: str, end_date: str, team_ids: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_TEAM_SPEND, + ENTITY_FIELD_TEAM, + _parse_csv_ids(team_ids), + start_date, + end_date, + ) + return resp.model_dump(mode="json") + + +async def _fetch_tag_usage_data( + start_date: str, end_date: str, tags: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_TAG_SPEND, + ENTITY_FIELD_TAG, + _parse_csv_ids(tags), + start_date, + end_date, + ) + return resp.model_dump(mode="json") + + +# --------------------------------------------------------------------------- +# Summarisers — convert raw JSON to concise text the LLM can reason over +# --------------------------------------------------------------------------- + + +def _accumulate_breakdown( + results: List[Dict[str, Any]], dimension: str, fields: List[str] +) -> Dict[str, Dict[str, float]]: + """Aggregate a single breakdown dimension across days.""" + totals: Dict[str, Dict[str, float]] = {} + for day in results: + for key, entry in day.get("breakdown", {}).get(dimension, {}).items(): + if key not in totals: + totals[key] = {f: 0.0 for f in fields} + m = entry.get("metrics", {}) + for f in fields: + totals[key][f] += m.get(f, 0) + return totals + + +def _ranked_lines( + totals: Dict[str, Dict[str, float]], + fmt: Callable[[str, Dict[str, float]], str], + limit: int, +) -> List[str]: + """Sort by spend descending, format each entry, and truncate.""" + return [ + fmt(name, vals) + for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[ + :limit + ] + ] + + +def _summarise_usage_data(data: Dict[str, Any]) -> str: + meta = data.get("metadata", {}) + results = data.get("results", []) + + header = ( + f"Total Spend: ${meta.get('total_spend', 0):.4f}\n" + f"Total Requests: {meta.get('total_api_requests', 0)}\n" + f"Successful: {meta.get('total_successful_requests', 0)} | " + f"Failed: {meta.get('total_failed_requests', 0)}\n" + f"Total Tokens: {meta.get('total_tokens', 0)}" + ) + + models = _accumulate_breakdown( + results, "models", ["spend", "api_requests", "total_tokens"] + ) + providers = _accumulate_breakdown(results, "providers", ["spend", "api_requests"]) + + model_lines = _ranked_lines( + models, + lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)", + TOP_N_MODELS, + ) + provider_lines = _ranked_lines( + providers, + lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs)", + TOP_N_PROVIDERS, + ) + + sections = [header, ""] + sections += ["Top Models by Spend:"] + (model_lines or [" (no data)"]) + [""] + sections += ["Top Providers by Spend:"] + (provider_lines or [" (no data)"]) + return "\n".join(sections) + + +def _summarise_entity_data(data: Dict[str, Any], entity_label: str) -> str: + """Summarise team/tag entity usage data.""" + results = data.get("results", []) + if not results: + return f"No {entity_label} usage data found for the given date range." + + totals: Dict[str, Dict[str, Any]] = {} + for day in results: + for eid, entry in day.get("breakdown", {}).get("entities", {}).items(): + if eid not in totals: + alias = entry.get("metadata", {}).get("alias", eid) + totals[eid] = {"alias": alias, "spend": 0.0, "requests": 0, "tokens": 0} + m = entry.get("metrics", {}) + totals[eid]["spend"] += m.get("spend", 0) + totals[eid]["requests"] += m.get("api_requests", 0) + totals[eid]["tokens"] += m.get("total_tokens", 0) + + lines = [f"{entity_label} Usage ({len(totals)} {entity_label.lower()}s):", ""] + for eid, d in sorted(totals.items(), key=lambda x: -x[1]["spend"]): + label = d["alias"] if d["alias"] != eid else eid + lines.append( + f"- {label} (ID: {eid}): ${d['spend']:.4f} | " + f"{int(d['requests'])} reqs | {int(d['tokens'])} tokens" + ) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Tool dispatch registry +# --------------------------------------------------------------------------- + +TOOL_HANDLERS: Dict[str, ToolHandler] = { + "get_usage_data": ToolHandler( + fetch=_fetch_usage_data, + summarise=_summarise_usage_data, + label="global usage data", + ), + "get_team_usage_data": ToolHandler( + fetch=_fetch_team_usage_data, + summarise=lambda data: _summarise_entity_data(data, "Team"), + label="team usage data", + ), + "get_tag_usage_data": ToolHandler( + fetch=_fetch_tag_usage_data, + summarise=lambda data: _summarise_entity_data(data, "Tag"), + label="tag usage data", + ), +} + + +# --------------------------------------------------------------------------- +# SSE streaming +# --------------------------------------------------------------------------- + + +def _sse(event: SSEEvent) -> str: + return f"data: {json.dumps(event)}\n\n" + + +def _resolve_fetch_kwargs( + fn_name: str, + fn_args: Dict[str, str], + user_id: Optional[str], + is_admin: bool, +) -> Dict[str, Any]: + """Build keyword arguments for a tool's fetch function.""" + start_date = fn_args.get("start_date", "") + end_date = fn_args.get("end_date", "") + if not start_date or not end_date: + raise ValueError("Missing required start_date or end_date from tool arguments") + kwargs: Dict[str, Any] = {"start_date": start_date, "end_date": end_date} + if fn_name == "get_usage_data": + if not is_admin: + kwargs["user_id"] = user_id + elif fn_args.get("user_id"): + kwargs["user_id"] = fn_args["user_id"] + elif fn_name == "get_team_usage_data" and fn_args.get("team_ids"): + kwargs["team_ids"] = fn_args["team_ids"] + elif fn_name == "get_tag_usage_data" and fn_args.get("tags"): + kwargs["tags"] = fn_args["tags"] + return kwargs + + +async def _execute_tool_call( + handler: ToolHandler, + fn_name: str, + fn_args: Dict[str, str], + user_id: Optional[str], + is_admin: bool, +) -> str: + """Run a single tool and return the summarised result text.""" + kwargs = _resolve_fetch_kwargs(fn_name, fn_args, user_id, is_admin) + raw_data = await handler["fetch"](**kwargs) + return handler["summarise"](raw_data) + + +async def _process_tool_call( + tc: Any, + chat_messages: List[Dict[str, Any]], + user_id: Optional[str], + is_admin: bool, +) -> AsyncIterator[str]: + """Execute a single tool call, yielding SSE events for status.""" + fn_name = tc.function.name + fn_args = json.loads(tc.function.arguments) + + allowed_names = {t["function"]["name"] for t in get_tools_for_role(is_admin)} + handler = TOOL_HANDLERS.get(fn_name) + + if fn_name not in allowed_names or not handler: + chat_messages.append( + { + "role": "tool", + "tool_call_id": tc.id, + "content": f"Tool not available: {fn_name}", + } + ) + return + + tool_event_base = { + "type": "tool_call", + "tool_name": fn_name, + "tool_label": handler["label"], + "arguments": fn_args, + } + yield _sse({**tool_event_base, "status": "running"}) + + try: + tool_result = await _execute_tool_call( + handler, fn_name, fn_args, user_id, is_admin + ) + yield _sse({**tool_event_base, "status": "complete"}) + except Exception as e: + verbose_proxy_logger.error("Tool %s failed: %s", fn_name, e) + tool_result = f"Error fetching {handler['label']}. Please try again." + yield _sse({**tool_event_base, "status": "error"}) + + chat_messages.append( + {"role": "tool", "tool_call_id": tc.id, "content": tool_result} + ) + + +async def _stream_final_response( + model: str, chat_messages: List[Dict[str, Any]] +) -> AsyncIterator[str]: + """Stream the final LLM response after tool results are appended.""" + yield _sse({"type": "status", "message": "Analyzing results..."}) + + response = await litellm.acompletion( + model=model, + messages=chat_messages, + stream=True, + temperature=USAGE_AI_TEMPERATURE, + ) + async for chunk in response: + delta = chunk.choices[0].delta.content + if delta: + yield _sse({"type": "chunk", "content": delta}) + + +async def stream_usage_ai_chat( + messages: List[Dict[str, str]], + model: Optional[str] = None, + user_id: Optional[str] = None, + is_admin: bool = False, +) -> AsyncIterator[str]: + """Stream SSE events: status → tool_call → chunk → done.""" + resolved_model = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL + truncated = ( + messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages + ) + chat_messages: List[Dict[str, Any]] = [ + {"role": "system", "content": _build_system_prompt(is_admin)}, + *truncated, + ] + + try: + yield _sse({"type": "status", "message": "Thinking..."}) + tools = get_tools_for_role(is_admin) + response = await litellm.acompletion( + model=resolved_model, + messages=chat_messages, + tools=tools, + temperature=USAGE_AI_TEMPERATURE, + ) + choice = response.choices[0] # type: ignore + + if not choice.message.tool_calls: + if choice.message.content: + yield _sse({"type": "chunk", "content": choice.message.content}) + yield _sse({"type": "done"}) + return + + chat_messages.append(choice.message.model_dump()) + for tc in choice.message.tool_calls: + async for event in _process_tool_call(tc, chat_messages, user_id, is_admin): + yield event + async for event in _stream_final_response(resolved_model, chat_messages): + yield event + yield _sse({"type": "done"}) + + except Exception as e: + verbose_proxy_logger.error("AI usage chat failed: %s", e) + yield _sse( + { + "type": "error", + "message": "An internal error occurred. Please try again.", + } + ) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py new file mode 100644 index 00000000000..0dbe518afb7 --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py @@ -0,0 +1,65 @@ +""" +USAGE AI CHAT ENDPOINTS + +/usage/ai/chat - Stream AI chat responses about usage data +""" + +from typing import List, Literal, Optional + +from fastapi import APIRouter, Depends, Request +from fastapi.responses import StreamingResponse +from pydantic import BaseModel, Field + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + + +class ChatMessage(BaseModel): + role: Literal["user", "assistant"] + content: str + + +class UsageAIChatRequest(BaseModel): + messages: List[ChatMessage] = Field( + ..., description="Chat messages (user/assistant history)" + ) + model: Optional[str] = Field(default=None, description="Model to use for AI chat") + + +@router.post( + "/usage/ai/chat", + tags=["Budget & Spend Tracking"], + dependencies=[Depends(user_api_key_auth)], +) +async def usage_ai_chat( + data: UsageAIChatRequest, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + AI chat about usage data. Streams SSE events with the AI response. + The AI agent has access to tools that query aggregated daily activity data. + """ + from litellm.proxy.management_endpoints.common_utils import ( + _user_has_admin_view, + ) + from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( + stream_usage_ai_chat, + ) + + is_admin = _user_has_admin_view(user_api_key_dict) + user_id = user_api_key_dict.user_id + messages = [{"role": m.role, "content": m.content} for m in data.messages] + + return StreamingResponse( + stream_usage_ai_chat( + messages=messages, + model=data.model, + user_id=user_id, + is_admin=is_admin, + ), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 36b6bcc7707..607306f3806 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -392,6 +392,7 @@ from litellm.proxy.management_endpoints.organization_endpoints import ( router as organization_router, ) from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router +from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router from litellm.proxy.management_endpoints.project_endpoints import ( router as project_router, ) @@ -12872,6 +12873,7 @@ app.include_router(caching_router) app.include_router(analytics_router) app.include_router(guardrails_router) app.include_router(policy_router) +app.include_router(usage_ai_router) app.include_router(policy_crud_router) app.include_router(policy_resolve_router) app.include_router(search_tool_management_router) diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py new file mode 100644 index 00000000000..f9303bd13a6 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -0,0 +1,402 @@ +""" +Tests for AI Usage Chat module. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( + TOOL_HANDLERS, + TOOLS_ADMIN, + TOOLS_BASE, + _build_system_prompt, + _summarise_entity_data, + _summarise_usage_data, + stream_usage_ai_chat, +) + + +SAMPLE_AGGREGATED_RESPONSE = { + "results": [ + { + "date": "2025-01-15", + "metrics": { + "spend": 50.25, + "prompt_tokens": 20000, + "completion_tokens": 10000, + "total_tokens": 30000, + "api_requests": 500, + "successful_requests": 480, + "failed_requests": 20, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + }, + "breakdown": { + "models": { + "gpt-4": { + "metrics": { + "spend": 40.0, + "api_requests": 300, + "total_tokens": 25000, + }, + "metadata": {}, + "api_key_breakdown": {}, + }, + }, + "providers": { + "openai": { + "metrics": {"spend": 50.25, "api_requests": 500}, + "metadata": {}, + "api_key_breakdown": {}, + }, + }, + "api_keys": { + "sk-test123": { + "metrics": {"spend": 50.25}, + "metadata": {"key_alias": "Production Key"}, + }, + }, + "model_groups": {}, + "mcp_servers": {}, + "entities": {}, + }, + }, + ], + "metadata": { + "total_spend": 50.25, + "total_api_requests": 500, + "total_successful_requests": 480, + "total_failed_requests": 20, + "total_tokens": 30000, + }, +} + +SAMPLE_TEAM_RESPONSE = { + "results": [ + { + "date": "2025-01-15", + "metrics": {"spend": 100.0, "api_requests": 1000, "total_tokens": 50000}, + "breakdown": { + "entities": { + "team-1": { + "metrics": { + "spend": 60.0, + "api_requests": 600, + "total_tokens": 30000, + }, + "metadata": {"alias": "Engineering"}, + "api_key_breakdown": {}, + }, + "team-2": { + "metrics": { + "spend": 40.0, + "api_requests": 400, + "total_tokens": 20000, + }, + "metadata": {"alias": "Marketing"}, + "api_key_breakdown": {}, + }, + }, + "models": {}, + "providers": {}, + "api_keys": {}, + "model_groups": {}, + "mcp_servers": {}, + }, + }, + ], + "metadata": {"total_spend": 100.0, "total_api_requests": 1000}, +} + + +class TestToolSchemas: + def test_admin_tools_include_all(self): + assert len(TOOLS_ADMIN) == 3 + names = {t["function"]["name"] for t in TOOLS_ADMIN} + assert "get_usage_data" in names + assert "get_team_usage_data" in names + assert "get_tag_usage_data" in names + + def test_base_tools_restricted_to_usage_only(self): + assert len(TOOLS_BASE) == 1 + assert TOOLS_BASE[0]["function"]["name"] == "get_usage_data" + + def test_admin_prompt_mentions_all_tools(self): + prompt = _build_system_prompt(is_admin=True) + assert "get_usage_data" in prompt + assert "get_team_usage_data" in prompt + assert "get_tag_usage_data" in prompt + + def test_non_admin_prompt_only_mentions_usage_tool(self): + prompt = _build_system_prompt(is_admin=False) + assert "get_usage_data" in prompt + assert "get_team_usage_data" not in prompt + assert "get_tag_usage_data" not in prompt + + def test_system_prompt_includes_todays_date(self): + from datetime import date + + prompt = _build_system_prompt(is_admin=True) + assert date.today().isoformat() in prompt + + +class TestSummariseUsageData: + def test_summarise_includes_totals(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "$50.25" in summary + assert "500" in summary + + def test_summarise_includes_models(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "gpt-4" in summary + + def test_summarise_includes_providers(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "openai" in summary + + def test_summarise_handles_empty_data(self): + empty = {"results": [], "metadata": {}} + summary = _summarise_usage_data(empty) + assert "no data" in summary.lower() + + +class TestSummariseEntityData: + def test_team_summary_includes_teams(self): + summary = _summarise_entity_data(SAMPLE_TEAM_RESPONSE, "Team") + assert "Engineering" in summary + assert "Marketing" in summary + assert "$60.0" in summary + assert "$40.0" in summary + + def test_team_summary_empty(self): + empty = {"results": [], "metadata": {}} + summary = _summarise_entity_data(empty, "Team") + assert "No Team usage data" in summary + + +class TestStreamUsageAiChat: + @pytest.mark.asyncio + async def test_stream_emits_status_events(self): + mock_tool_call = MagicMock() + mock_tool_call.id = "call_123" + mock_tool_call.function.name = "get_usage_data" + mock_tool_call.function.arguments = json.dumps( + { + "start_date": "2025-01-01", + "end_date": "2025-01-31", + } + ) + + mock_first_response = MagicMock() + mock_first_response.choices = [MagicMock()] + mock_first_response.choices[0].message.tool_calls = [mock_tool_call] + mock_first_response.choices[0].message.model_dump.return_value = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_usage_data", + "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}', + }, + } + ], + } + + async def mock_stream(): + chunk = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Total spend is $50.25" + yield chunk + + with patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm, patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_usage_data", + new_callable=AsyncMock, + ) as mock_fetch: + mock_litellm.acompletion = AsyncMock( + side_effect=[ + mock_first_response, + mock_stream(), + ] + ) + mock_fetch.return_value = SAMPLE_AGGREGATED_RESPONSE + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "What is my total spend?"}], + model="gpt-4o-mini", + user_id="user-123", + is_admin=True, + ): + events.append(json.loads(event.replace("data: ", "").strip())) + + status_events = [e for e in events if e["type"] == "status"] + tool_call_events = [e for e in events if e["type"] == "tool_call"] + chunk_events = [e for e in events if e["type"] == "chunk"] + done_events = [e for e in events if e["type"] == "done"] + + assert len(status_events) >= 1 + assert "Thinking" in status_events[0]["message"] + assert len(tool_call_events) >= 1 + assert tool_call_events[0]["tool_name"] == "get_usage_data" + assert tool_call_events[0]["status"] in ("running", "complete") + assert len(chunk_events) >= 1 + assert len(done_events) == 1 + + @pytest.mark.asyncio + async def test_stream_handles_team_tool(self): + mock_tool_call = MagicMock() + mock_tool_call.id = "call_team" + mock_tool_call.function.name = "get_team_usage_data" + mock_tool_call.function.arguments = json.dumps( + { + "start_date": "2025-01-01", + "end_date": "2025-01-31", + } + ) + + mock_first_response = MagicMock() + mock_first_response.choices = [MagicMock()] + mock_first_response.choices[0].message.tool_calls = [mock_tool_call] + mock_first_response.choices[0].message.model_dump.return_value = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_team", + "type": "function", + "function": { + "name": "get_team_usage_data", + "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}', + }, + } + ], + } + + async def mock_stream(): + chunk = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Engineering is the top team." + yield chunk + + with patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm, patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_team_usage_data", + new_callable=AsyncMock, + ) as mock_fetch: + mock_litellm.acompletion = AsyncMock( + side_effect=[ + mock_first_response, + mock_stream(), + ] + ) + mock_fetch.return_value = SAMPLE_TEAM_RESPONSE + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "Which team spends the most?"}], + model="gpt-4o-mini", + is_admin=True, + ): + events.append(json.loads(event.replace("data: ", "").strip())) + + chunk_events = [e for e in events if e["type"] == "chunk"] + assert len(chunk_events) >= 1 + assert "Engineering" in chunk_events[0]["content"] + + @pytest.mark.asyncio + async def test_stream_handles_error(self): + with patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm: + mock_litellm.acompletion = AsyncMock(side_effect=Exception("LLM error")) + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "test"}], + ): + events.append(json.loads(event.replace("data: ", "").strip())) + + error_events = [e for e in events if e["type"] == "error"] + assert len(error_events) == 1 + assert "internal error" in error_events[0]["message"].lower() + + @pytest.mark.asyncio + async def test_non_admin_enforces_user_id(self): + mock_tool_call = MagicMock() + mock_tool_call.id = "call_456" + mock_tool_call.function.name = "get_usage_data" + mock_tool_call.function.arguments = json.dumps( + { + "start_date": "2025-01-01", + "end_date": "2025-01-31", + "user_id": "other-user", + } + ) + + mock_first_response = MagicMock() + mock_first_response.choices = [MagicMock()] + mock_first_response.choices[0].message.tool_calls = [mock_tool_call] + mock_first_response.choices[0].message.model_dump.return_value = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_456", + "type": "function", + "function": { + "name": "get_usage_data", + "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31","user_id":"other-user"}', + }, + } + ], + } + + async def mock_stream(): + chunk = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Data." + yield chunk + + mock_fetch = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE) + + with patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm, patch.dict( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.TOOL_HANDLERS", + { + "get_usage_data": { + "fetch": mock_fetch, + "summarise": _summarise_usage_data, + "label": "global usage data", + } + }, + ): + mock_litellm.acompletion = AsyncMock( + side_effect=[ + mock_first_response, + mock_stream(), + ] + ) + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "Show data"}], + model="gpt-4o-mini", + user_id="my-user-id", + is_admin=False, + ): + events.append(event) + + mock_fetch.assert_called_once_with( + start_date="2025-01-01", + end_date="2025-01-31", + user_id="my-user-id", + ) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 489a39a7ee2..3787f451ad3 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -1757,29 +1757,6 @@ "url": "https://opencollective.com/libvips" } }, - "node_modules/@isaacs/balanced-match": { - "version": "4.0.1", - "resolved": "https://registry.npmjs.org/@isaacs/balanced-match/-/balanced-match-4.0.1.tgz", - "integrity": "sha512-yzMTt9lEb8Gv7zRioUilSglI0c0smZ9k5D65677DLWLtWJaXIS3CqcGyUFByYKlnUj6TkjLVs54fBl6+TiGQDQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": "20 || >=22" - } - }, - "node_modules/@isaacs/brace-expansion": { - "version": "5.0.1", - "resolved": "https://registry.npmjs.org/@isaacs/brace-expansion/-/brace-expansion-5.0.1.tgz", - "integrity": "sha512-WMz71T1JS624nWj2n2fnYAuPovhv7EUhk69R6i9dsVyzxt5eM3bjwvgk9L+APE1TRscGysAVMANkB0jh0LQZrQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "@isaacs/balanced-match": "^4.0.1" - }, - "engines": { - "node": "20 || >=22" - } - }, "node_modules/@istanbuljs/schema": { "version": "0.1.3", "resolved": "https://registry.npmjs.org/@istanbuljs/schema/-/schema-0.1.3.tgz", @@ -3696,32 +3673,6 @@ "typescript": ">=4.8.4 <6.0.0" } }, - "node_modules/@typescript-eslint/typescript-estree/node_modules/brace-expansion": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.2.tgz", - "integrity": "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "balanced-match": "^1.0.0" - } - }, - "node_modules/@typescript-eslint/typescript-estree/node_modules/minimatch": { - "version": "9.0.5", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-9.0.5.tgz", - "integrity": "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==", - "dev": true, - "license": "ISC", - "dependencies": { - "brace-expansion": "^2.0.1" - }, - "engines": { - "node": ">=16 || 14 >=14.17" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/@typescript-eslint/utils": { "version": "8.54.0", "resolved": "https://registry.npmjs.org/@typescript-eslint/utils/-/utils-8.54.0.tgz", @@ -4749,11 +4700,14 @@ } }, "node_modules/balanced-match": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz", - "integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw==", + "version": "4.0.4", + "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-4.0.4.tgz", + "integrity": "sha512-BLrgEcRTwX2o6gGxGOCNyMvGSp35YofuYzw9h1IMTRmKqttAZZVU67bdb9Pr2vUHA8+j3i2tJfjO6C6+4myGTA==", "dev": true, - "license": "MIT" + "license": "MIT", + "engines": { + "node": "18 || 20 || >=22" + } }, "node_modules/baseline-browser-mapping": { "version": "2.9.19", @@ -4787,14 +4741,16 @@ } }, "node_modules/brace-expansion": { - "version": "1.1.12", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.12.tgz", - "integrity": "sha512-9T9UjW3r0UW5c1Q7GTwllptXwhvYmEzFhzMfZ9H7FQWt+uZePjZPjBP/W1ZEyZ1twGWom5/56TF4lPcqjnDHcg==", + "version": "5.0.3", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.3.tgz", + "integrity": "sha512-fy6KJm2RawA5RcHkLa1z/ScpBeA762UF9KmZQxwIbDtRJrgLzM10depAiEQ+CXYcoiqW1/m96OAAoke2nE9EeA==", "dev": true, "license": "MIT", "dependencies": { - "balanced-match": "^1.0.0", - "concat-map": "0.0.1" + "balanced-match": "^4.0.2" + }, + "engines": { + "node": "18 || 20 || >=22" } }, "node_modules/braces": { @@ -5149,13 +5105,6 @@ "integrity": "sha512-VRhuHOLoKYOy4UbilLbUzbYg93XLjv2PncJC50EuTWPA3gaja1UjBsUP/D/9/juV3vQFr6XBEzn9KCAHdUvOHw==", "license": "MIT" }, - "node_modules/concat-map": { - "version": "0.0.1", - "resolved": "https://registry.npmjs.org/concat-map/-/concat-map-0.0.1.tgz", - "integrity": "sha512-/Srv4dswyQNBfohGpz9o6Yb3Gz3SrUDqBH5rTuhGR7ahtlbYKnVxw2bCFMRljaA7EXHaXZ8wsHdodFvbkhKmqg==", - "dev": true, - "license": "MIT" - }, "node_modules/copy-to-clipboard": { "version": "3.3.3", "resolved": "https://registry.npmjs.org/copy-to-clipboard/-/copy-to-clipboard-3.3.3.tgz", @@ -6924,22 +6873,6 @@ "node": ">=10.13.0" } }, - "node_modules/glob/node_modules/minimatch": { - "version": "10.1.1", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.1.1.tgz", - "integrity": "sha512-enIvLvRAFZYXJzkCYG5RKmPfrFArdLv+R+lbQ53BmIMLIry74bjKzX6iHAm8WYamJkhSSEabrWN5D97XnKObjQ==", - "dev": true, - "license": "BlueOak-1.0.0", - "dependencies": { - "@isaacs/brace-expansion": "^5.0.0" - }, - "engines": { - "node": "20 || >=22" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/globals": { "version": "14.0.0", "resolved": "https://registry.npmjs.org/globals/-/globals-14.0.0.tgz", @@ -9035,16 +8968,19 @@ } }, "node_modules/minimatch": { - "version": "3.1.2", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz", - "integrity": "sha512-J7p63hRiAjw1NDEww1W7i37+ByIrOWO5XQQAzZ3VOcL0PNybwpfmV/N05zFAzwQ9USyEcX6t3UO+K5aqBQOIHw==", + "version": "10.2.2", + "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.2.tgz", + "integrity": "sha512-+G4CpNBxa5MprY+04MbgOw1v7So6n5JY166pFi9KfYwT78fxScCeSNQSNzp6dpPSW2rONOps6Ocam1wFhCgoVw==", "dev": true, - "license": "ISC", + "license": "BlueOak-1.0.0", "dependencies": { - "brace-expansion": "^1.1.7" + "brace-expansion": "^5.0.2" }, "engines": { - "node": "*" + "node": "18 || 20 || >=22" + }, + "funding": { + "url": "https://github.com/sponsors/isaacs" } }, "node_modules/minimist": { @@ -12004,32 +11940,6 @@ "node": ">=18" } }, - "node_modules/test-exclude/node_modules/brace-expansion": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.2.tgz", - "integrity": "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "balanced-match": "^1.0.0" - } - }, - "node_modules/test-exclude/node_modules/minimatch": { - "version": "9.0.5", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-9.0.5.tgz", - "integrity": "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==", - "dev": true, - "license": "ISC", - "dependencies": { - "brace-expansion": "^2.0.1" - }, - "engines": { - "node": ">=16 || 14 >=14.17" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/thenify": { "version": "3.3.1", "resolved": "https://registry.npmjs.org/thenify/-/thenify-3.3.1.tgz", @@ -13085,21 +12995,6 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } - }, - "node_modules/@next/swc-win32-ia32-msvc": { - "version": "14.2.33", - "resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz", - "integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==", - "cpu": [ - "ia32" - ], - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 10" - } } } } diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx new file mode 100644 index 00000000000..85ce11e605a --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx @@ -0,0 +1,85 @@ +import { screen } from "@testing-library/react"; +import { beforeAll, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; +import UsageAIChatPanel from "./UsageAIChatPanel"; + +beforeAll(() => { + if (typeof window !== "undefined" && !window.ResizeObserver) { + window.ResizeObserver = class ResizeObserver { + observe() {} + unobserve() {} + disconnect() {} + } as any; + } +}); + +vi.mock("../../networking", () => ({ + modelHubCall: vi.fn().mockResolvedValue({ + data: [ + { model_group: "gpt-4" }, + { model_group: "claude-3-opus" }, + ], + }), + usageAiChatStream: vi.fn(), +})); + +const defaultProps = { + open: true, + onClose: vi.fn(), + accessToken: "test-token", +}; + +describe("UsageAIChatPanel", () => { + it("should render the panel when open", () => { + renderWithProviders(); + + expect(screen.getByText("Ask AI")).toBeInTheDocument(); + expect( + screen.getByText("Ask about your spend, models, keys, and trends") + ).toBeInTheDocument(); + }); + + it("should render model selector", () => { + renderWithProviders(); + + expect(screen.getByText("Select a model (optional, defaults to gpt-4o-mini)")).toBeInTheDocument(); + }); + + it("should render empty state message when no conversation", () => { + renderWithProviders(); + + expect(screen.getByText("Ask a question about your usage")).toBeInTheDocument(); + }); + + it("should render the send button", () => { + renderWithProviders(); + + expect(screen.getByText("Send")).toBeInTheDocument(); + }); + + it("should render input placeholder", () => { + renderWithProviders(); + + expect(screen.getByPlaceholderText("Ask about your usage...")).toBeInTheDocument(); + }); + + it("should render clear chat button", () => { + renderWithProviders(); + + expect(screen.getByText("Clear chat")).toBeInTheDocument(); + }); + + it("should have the panel element even when closed (just off-screen)", () => { + renderWithProviders(); + + expect(screen.getByTestId("usage-ai-chat-panel")).toBeInTheDocument(); + expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-full"); + }); + + it("should not have translate-x-full class when open", () => { + renderWithProviders(); + + expect(screen.getByTestId("usage-ai-chat-panel")).not.toHaveClass("translate-x-full"); + expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-0"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx new file mode 100644 index 00000000000..85f46bfa346 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx @@ -0,0 +1,402 @@ +import React, { useEffect, useRef, useState } from "react"; +import { Button, Select, Input, Spin } from "antd"; +import ReactMarkdown from "react-markdown"; +import { modelHubCall, usageAiChatStream, UsageAiToolCallEvent } from "../../networking"; + +const { TextArea } = Input; + +interface ToolCallStep { + tool_name: string; + tool_label: string; + arguments: Record; + status: "running" | "complete" | "error"; + error?: string; +} + +interface ChatMessage { + role: "user" | "assistant"; + content: string; + toolCalls?: ToolCallStep[]; +} + +interface UsageAIChatPanelProps { + open: boolean; + onClose: () => void; + accessToken: string | null; +} + +const TOOL_ICONS: Record = { + get_usage_data: "📊", + get_team_usage_data: "👥", + get_tag_usage_data: "🏷️", +}; + +const ToolCallDisplay: React.FC<{ step: ToolCallStep }> = ({ step }) => { + const icon = TOOL_ICONS[step.tool_name] || "🔧"; + const args = step.arguments; + const dateRange = args.start_date && args.end_date + ? `${args.start_date} → ${args.end_date}` + : ""; + const filter = args.team_ids || args.tags || args.user_id || ""; + + return ( +
+ + {step.status === "running" ? ( + + ) : step.status === "error" ? ( + ✗ + ) : ( + ✓ + )} + +
+
+ {icon} {step.tool_label} +
+ {dateRange && ( +
{dateRange}
+ )} + {filter && ( +
Filter: {filter}
+ )} + {step.status === "error" && step.error && ( +
{step.error}
+ )} +
+
+ ); +}; + +const MarkdownContent: React.FC<{ content: string }> = ({ content }) => ( +

{children}

, + strong: ({ children }) => {children}, + ul: ({ children }) =>
    {children}
, + ol: ({ children }) =>
    {children}
, + li: ({ children }) =>
  • {children}
    • , + h1: ({ children }) =>

    {children}

    , + h2: ({ children }) =>

    {children}

    , + h3: ({ children }) =>

    {children}

    , + code: ({ children, className }) => { + const isBlock = className?.includes("language-"); + return isBlock ? ( + 
    +            {children}
+
    + ) : ( + {children} + ); + }, + table: ({ children }) => ( +
    + {children}
    +
    + ), + th: ({ children }) => {children}, + td: ({ children }) => {children}, + }} + > + {content} +     +); + +const UsageAIChatPanel: React.FC = ({ + open, + onClose, + accessToken, +}) => { + const [messages, setMessages] = useState([]); + const [inputText, setInputText] = useState(""); + const [isLoading, setIsLoading] = useState(false); + const [selectedModel, setSelectedModel] = useState(undefined); + const [availableModels, setAvailableModels] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + const [streamingContent, setStreamingContent] = useState(""); + const [statusMessage, setStatusMessage] = useState(null); + const [activeToolCalls, setActiveToolCalls] = useState([]); + const messagesEndRef = useRef(null); + const abortControllerRef = useRef(null); + + useEffect(() => { + if (open && availableModels.length === 0) { + loadModels(); + } + }, [open]); + + useEffect(() => { + if (typeof messagesEndRef.current?.scrollIntoView === "function") { + messagesEndRef.current.scrollIntoView({ behavior: "smooth" }); + } + }, [messages, streamingContent, activeToolCalls, statusMessage]); + + const loadModels = async () => { + if (!accessToken) return; + setIsLoadingModels(true); + try { + const fetchedModels = await modelHubCall(accessToken); + if (fetchedModels?.data?.length > 0) { + const models = fetchedModels.data + .map((item: any) => item.model_group as string) + .sort(); + setAvailableModels(models); + } + } catch (error) { + console.error("Failed to load models:", error); + } finally { + setIsLoadingModels(false); + } + }; + + const handleSend = async () => { + if (!accessToken || !inputText.trim() || isLoading) return; + + const userMessage: ChatMessage = { role: "user", content: inputText.trim() }; + const updatedMessages = [...messages, userMessage]; + setMessages(updatedMessages); + setInputText(""); + setIsLoading(true); + setStreamingContent(""); + setStatusMessage(null); + setActiveToolCalls([]); + + const abortController = new AbortController(); + abortControllerRef.current = abortController; + + let accumulated = ""; + const toolCalls: ToolCallStep[] = []; + + try { + await usageAiChatStream( + accessToken, + updatedMessages.slice(-20).map((m) => ({ role: m.role, content: m.content })), + selectedModel || "", + (content: string) => { + setStatusMessage(null); + accumulated += content; + setStreamingContent(accumulated); + }, + () => { + setStatusMessage(null); + setActiveToolCalls([]); + setMessages((prev) => [ + ...prev, + { role: "assistant", content: accumulated, toolCalls: toolCalls.length > 0 ? [...toolCalls] : undefined }, + ]); + setStreamingContent(""); + }, + (errorMsg: string) => { + setStatusMessage(null); + setActiveToolCalls([]); + setMessages((prev) => [ + ...prev, + { role: "assistant", content: `Error: ${errorMsg}` }, + ]); + setStreamingContent(""); + }, + (status: string) => { + setStatusMessage(status); + }, + (event: UsageAiToolCallEvent) => { + const idx = toolCalls.findIndex((tc) => tc.tool_name === event.tool_name); + if (idx >= 0) { + toolCalls[idx] = { ...event }; + } else { + toolCalls.push({ ...event }); + } + setActiveToolCalls([...toolCalls]); + }, + abortController.signal, + ); + } catch (error: any) { + if (error?.name === "AbortError" || abortController.signal.aborted) { + return; + } + const errorMsg = error?.message || "Failed to get response. Please try again."; + setMessages((prev) => [ + ...prev, + { role: "assistant", content: `Error: ${errorMsg}` }, + ]); + setStreamingContent(""); + } finally { + setIsLoading(false); + abortControllerRef.current = null; + } + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + handleSend(); + } + }; + + const handleClose = () => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + } + onClose(); + }; + + const handleClear = () => { + setMessages([]); + setStreamingContent(""); + setActiveToolCalls([]); + setStatusMessage(null); + }; + + return ( +
    + {/* Header */} +
    +
    +
    + + + +

    Ask AI

    +
    + +
    +

    + Ask about your spend, models, keys, and trends +

    +
    + + {/* Model selector */} +
    +