From 219cd414f39993c85cbd66c31307758e1b8394f6 Mon Sep 17 00:00:00 2001 From: Chesars Date: Wed, 21 Jan 2026 14:54:27 -0300 Subject: [PATCH 01/58] Fix model management page UI improvements - Delete button now shows confirmation modal instead of redirecting to detail page - Entire table row is clickable to view model details - Fixed Last Refreshed text alignment and shortened time format - Added dismiss button (X) to Missing provider banner with localStorage persistence - Added compact Request Provider button in header when banner is dismissed --- .../ModelsAndEndpointsView.tsx | 116 ++++++++++++------ .../components/AllModelsTab.tsx | 65 +++++++++- .../src/components/model_dashboard/table.tsx | 8 +- .../components/molecules/models/columns.tsx | 18 ++- 4 files changed, 160 insertions(+), 47 deletions(-) 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 6a4882a92a2..9a96cb33a5c 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 @@ -15,7 +15,7 @@ import { transformModelData } from "./utils/modelDataTransformer"; import { all_admin_roles, internalUserRoles, isProxyAdminRole, isUserTeamAdminForAnyTeam } from "@/utils/roles"; import { RefreshIcon } from "@heroicons/react/outline"; import { useQueryClient } from "@tanstack/react-query"; -import { Col, Grid, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react"; +import { Col, Grid, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; import type { UploadProps } from "antd"; import { Form, Typography } from "antd"; import { PlusCircleOutlined } from "@ant-design/icons"; @@ -62,6 +62,12 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te const [selectedModelId, setSelectedModelId] = useState(null); const [selectedTeamId, setSelectedTeamId] = useState(null); const [selectedTabIndex, setSelectedTabIndex] = useState(0); + const [showMissingProviderBanner, setShowMissingProviderBanner] = useState(() => { + if (typeof window !== "undefined") { + return localStorage.getItem("hideMissingProviderBanner") !== "true"; + } + return true; + }); const queryClient = useQueryClient(); const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo(); @@ -153,7 +159,7 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te const handleRefreshClick = () => { const currentDate = new Date(); - setLastRefreshed(currentDate.toLocaleString()); + setLastRefreshed(currentDate.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })); queryClient.invalidateQueries({ queryKey: ["models", "list"] }); refetchModels(); }; @@ -275,43 +281,75 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te

Add and manage models for the proxy

)} + {!showMissingProviderBanner && ( + + + Request Provider + + )} {/* Missing Provider Banner */} -
-
- -
-
-

Missing a provider?

-

- The LiteLLM engineering team is constantly adding support for new LLM models, providers, endpoints. If - you don't see the one you need, let us know and we'll prioritize it. -

-
- - Request Provider - +
+ +
+
+

Missing a provider?

+

+ The LiteLLM engineering team is constantly adding support for new LLM models, providers, endpoints. If + you don't see the one you need, let us know and we'll prioritize it. +

+
+ - - - -
+ Request Provider + + + + + + + )} {selectedModelId && !isLoading ? ( = ({ premiumUser, te {all_admin_roles.includes(userRole) && Price Data Reload} -
- {lastRefreshed && Last Refreshed: {lastRefreshed}} +
+ {lastRefreshed && Last Refreshed: {lastRefreshed}}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index a5553fa5f56..25ca5f4cf5e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -5,8 +5,12 @@ import { Team } from "@/components/key_team_helpers/key_list"; import { AllModelsDataTable } from "@/components/model_dashboard/all_models_table"; import { columns } from "@/components/molecules/models/columns"; import { getDisplayModelName } from "@/components/view_model/model_name_display"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { modelDeleteCall } from "@/components/networking"; import { InfoCircleOutlined } from "@ant-design/icons"; import { PaginationState, SortingState } from "@tanstack/react-table"; +import { useQueryClient } from "@tanstack/react-query"; import { Grid, Select, SelectItem, TabPanel, Text } from "@tremor/react"; import { Skeleton, Spin } from "antd"; import debounce from "lodash/debounce"; @@ -33,8 +37,9 @@ const AllModelsTab = ({ setSelectedTeamId, }: AllModelsTabProps) => { const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap(); - const { userId, userRole, premiumUser } = useAuthorized(); + const { accessToken, userId, userRole, premiumUser } = useAuthorized(); const { data: teams, isLoading: isLoadingTeams } = useTeams(); + const queryClient = useQueryClient(); const [modelNameSearch, setModelNameSearch] = useState(""); const [debouncedSearch, setDebouncedSearch] = useState(""); @@ -92,7 +97,7 @@ const AllModelsTab = ({ return sort.desc ? "desc" : "asc"; }, [sorting]); - const { data: rawModelData, isLoading: isLoadingModelsInfo } = useModelsInfo( + const { data: rawModelData, isLoading: isLoadingModelsInfo, refetch: refetchModels } = useModelsInfo( currentPage, pageSize, debouncedSearch || undefined, @@ -117,6 +122,9 @@ const AllModelsTab = ({ return transformModelData(rawModelData, getProviderFromModel); }, [rawModelData, modelCostMapData]); + const [deleteModalModelId, setDeleteModalModelId] = useState(null); + const [deleteLoading, setDeleteLoading] = useState(false); + // Get pagination metadata from the response const paginationMeta = useMemo(() => { if (!rawModelData) { @@ -187,6 +195,28 @@ const AllModelsTab = ({ setSorting([]); }; + const modelToDelete = useMemo(() => { + if (!deleteModalModelId || !modelData?.data) return null; + return modelData.data.find((model: any) => model.model_info.id === deleteModalModelId); + }, [deleteModalModelId, modelData]); + + const handleDeleteModel = async () => { + if (!accessToken || !deleteModalModelId) return; + try { + setDeleteLoading(true); + await modelDeleteCall(accessToken, deleteModalModelId); + NotificationsManager.success("Model deleted successfully"); + queryClient.invalidateQueries({ queryKey: ["models", "list"] }); + refetchModels(); + } catch (error) { + console.error("Error deleting model:", error); + NotificationsManager.fromBackend("Failed to delete model"); + } finally { + setDeleteLoading(false); + setDeleteModalModelId(null); + } + }; + return ( @@ -481,6 +511,7 @@ const AllModelsTab = ({ () => { }, expandedRows, setExpandedRows, + setDeleteModalModelId, )} data={filteredData} isLoading={isLoadingModelsInfo} @@ -489,10 +520,40 @@ const AllModelsTab = ({ pagination={pagination} onPaginationChange={setPagination} enablePagination={true} + onRowClick={(model: any) => setSelectedModelId(model.model_info.id)} />
+ + setDeleteModalModelId(null)} + onOk={handleDeleteModel} + confirmLoading={deleteLoading} + /> ); }; diff --git a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx index 79224edba43..b7c416a2199 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx @@ -30,6 +30,7 @@ interface ModelDataTableProps { pagination?: PaginationState; onPaginationChange?: OnChangeFn; enablePagination?: boolean; + onRowClick?: (row: TData) => void; } export function ModelDataTable({ @@ -40,6 +41,7 @@ export function ModelDataTable({ pagination, onPaginationChange, enablePagination = false, + onRowClick, }: ModelDataTableProps) { const [sorting, setSorting] = React.useState(defaultSorting); const [columnResizeMode] = React.useState("onChange"); @@ -157,7 +159,11 @@ export function ModelDataTable({ ) : tableInstance.getRowModel().rows.length > 0 ? ( tableInstance.getRowModel().rows.map((row) => ( - + onRowClick?.(row.original)} + className={onRowClick ? "cursor-pointer hover:bg-gray-50" : ""} + > {row.getVisibleCells().map((cell) => ( void, expandedRows: Set, setExpandedRows: (expandedRows: Set) => void, + onDeleteClick?: (modelId: string) => void, ): ColumnDef[] => [ { header: () => Model ID, @@ -27,7 +28,10 @@ export const columns = (
setSelectedModelId(model.model_info.id)} + onClick={(e) => { + e.stopPropagation(); + setSelectedModelId(model.model_info.id); + }} > {model.model_info.id}
@@ -195,7 +199,10 @@ export const columns = ( size="xs" variant="light" className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]" - onClick={() => setSelectedTeamId(model.model_info.team_id)} + onClick={(e: React.MouseEvent) => { + e.stopPropagation(); + setSelectedTeamId(model.model_info.team_id); + }} > {model.model_info.team_id.slice(0, 7)}... @@ -300,9 +307,10 @@ export const columns = ( { - if (canEditModel) { - setSelectedModelId(model.model_info.id); + onClick={(e) => { + e.stopPropagation(); + if (canEditModel && onDeleteClick) { + onDeleteClick(model.model_info.id); } }} className={!canEditModel ? "opacity-50 cursor-not-allowed" : "cursor-pointer hover:text-red-600"} From c005f1aec21dd5b043c8a89705ad89cdebdcd58f Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 26 Jan 2026 12:17:13 -0300 Subject: [PATCH 02/58] test(ui): add tests for model management UI improvements - Add tests for dismissable Missing provider banner with localStorage persistence - Add tests for compact Request Provider button when banner is dismissed - Add tests for delete modal functionality and DB Model badge - Add tests for clickable Model ID that calls setSelectedModelId --- .../ModelsAndEndpointsView.test.tsx | 99 +++++++++++++- .../components/AllModelsTab.test.tsx | 121 +++++++++++++++++- 2 files changed, 218 insertions(+), 2 deletions(-) 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..f24de9f673c 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,9 +1,27 @@ /* @vitest-environment jsdom */ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render } from "@testing-library/react"; +import { fireEvent, render } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ModelsAndEndpointsView from "./ModelsAndEndpointsView"; +// Mock localStorage +const localStorageMock = (() => { + let store: Record = {}; + return { + getItem: (key: string) => store[key] || null, + setItem: (key: string, value: string) => { + store[key] = value; + }, + removeItem: (key: string) => { + delete store[key]; + }, + clear: () => { + store = {}; + }, + }; +})(); +Object.defineProperty(window, "localStorage", { value: localStorageMock }); + // Minimal stubs to avoid Next.js router and network usage during render vi.mock("@/components/networking", () => ({ credentialListCall: vi.fn().mockResolvedValue({ credentials: [] }), @@ -104,4 +122,83 @@ describe("ModelsAndEndpointsView", () => { ); expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument(); }, 15000); + + it("should show Missing provider banner by default", async () => { + localStorageMock.clear(); + const queryClient = createQueryClient(); + const { findByText } = render( + + {}} + premiumUser={false} + teams={[]} + /> + , + ); + expect(await findByText("Missing a provider?", {}, { timeout: 10000 })).toBeInTheDocument(); + }, 15000); + + it("should hide Missing provider banner when dismiss button is clicked and persist to localStorage", async () => { + localStorageMock.clear(); + const queryClient = createQueryClient(); + const { findByText, queryByText, container } = render( + + {}} + premiumUser={false} + teams={[]} + /> + , + ); + + // Wait for banner to appear + expect(await findByText("Missing a provider?", {}, { timeout: 10000 })).toBeInTheDocument(); + + // Find and click dismiss button (X button) + const dismissButton = container.querySelector('button[aria-label="Dismiss banner"]'); + if (dismissButton) { + fireEvent.click(dismissButton); + + // Banner should be hidden + expect(queryByText("Missing a provider?")).not.toBeInTheDocument(); + + // LocalStorage should be updated + expect(localStorageMock.getItem("hideMissingProviderBanner")).toBe("true"); + } + }, 15000); + + it("should show compact Request Provider button when banner is dismissed", async () => { + // Set localStorage to hide banner + localStorageMock.setItem("hideMissingProviderBanner", "true"); + const queryClient = createQueryClient(); + const { findByText, queryByText } = render( + + {}} + premiumUser={false} + teams={[]} + /> + , + ); + + // Wait for component to render + await findByText("Model Management", {}, { timeout: 10000 }); + + // Banner should not be visible + expect(queryByText("Missing a provider?")).not.toBeInTheDocument(); + + // Compact Request Provider button should be visible in header + const requestProviderLinks = document.querySelectorAll('a[href="https://models.litellm.ai/?request=true"]'); + // There should be a compact button when banner is hidden + expect(requestProviderLinks.length).toBeGreaterThan(0); + }, 15000); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index 813a365d367..30c39bd6a2c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -1,8 +1,30 @@ import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized"; -import { render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import AllModelsTab from "./AllModelsTab"; +// Mock modelDeleteCall +const mockModelDeleteCall = vi.fn().mockResolvedValue({}); +vi.mock("@/components/networking", () => ({ + modelDeleteCall: (...args: any[]) => mockModelDeleteCall(...args), +})); + +// Mock NotificationsManager +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +// Mock react-query +const mockInvalidateQueries = vi.fn(); +vi.mock("@tanstack/react-query", () => ({ + useQueryClient: () => ({ + invalidateQueries: mockInvalidateQueries, + }), +})); + // Mock the useModelsInfo hook const mockUseModelsInfo = vi.fn(() => ({ data: { data: [], total_count: 0, current_page: 1, total_pages: 1, size: 50 }, @@ -493,4 +515,101 @@ describe("AllModelsTab", () => { const previousButton = screen.getByRole("button", { name: /previous/i }); expect(previousButton).toBeDisabled(); }); + + it("should pass setDeleteModalModelId to columns for delete functionality", async () => { + // This test verifies that the delete modal setter is passed to columns + // The actual modal rendering is handled by DeleteResourceModal component + mockUseTeams.mockReturnValue({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), + }); + + mockUseModelCostMap.mockReturnValue( + createModelCostMapMock({ + "gpt-4-delete-test": { litellm_provider: "openai" }, + }), + ); + + const modelData = createPaginatedModelData([ + { + model_name: "gpt-4-delete-test", + litellm_model_name: "gpt-4-delete-test", + provider: "openai", + model_info: { + id: "model-to-delete", + db_model: true, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", + }, + }, + ], 1, 1, 1, 50); + + mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() }); + + render(); + + await waitFor(() => { + expect(screen.getByText("gpt-4-delete-test")).toBeInTheDocument(); + }); + + // Verify the DB Model badge is shown (indicating it can be deleted) + expect(screen.getByText("DB Model")).toBeInTheDocument(); + }); + + it("should render clickable model ID that calls setSelectedModelId", async () => { + mockUseTeams.mockReturnValue({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), + }); + + mockUseModelCostMap.mockReturnValue( + createModelCostMapMock({ + "gpt-4-clickable": { litellm_provider: "openai" }, + }), + ); + + const modelData = createPaginatedModelData([ + { + model_name: "gpt-4-clickable", + litellm_model_name: "gpt-4-clickable", + provider: "openai", + model_info: { + id: "clickable-model-id", + db_model: true, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", + }, + }, + ], 1, 1, 1, 50); + + mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() }); + + render(); + + await waitFor(() => { + expect(screen.getByText("gpt-4-clickable")).toBeInTheDocument(); + }); + + // Click on the Model ID cell which should call setSelectedModelId + const modelIdCell = screen.getByText("clickable-model-id"); + expect(modelIdCell).toBeInTheDocument(); + + fireEvent.click(modelIdCell); + + await waitFor(() => { + expect(mockSetSelectedModelId).toHaveBeenCalledWith("clickable-model-id"); + }); + }); }); From 5e24937ab53cb83611caea6b6c30eeafc076ab44 Mon Sep 17 00:00:00 2001 From: "Srikanth @adobe" Date: Tue, 3 Mar 2026 20:21:55 -0800 Subject: [PATCH 03/58] [Issue#21344]: avoid migration hook serviceaccount dependency cycle (#21405) * helm cyclic dependency fix * updating test case for handling edge cases --- .../litellm-helm/templates/_helpers.tpl | 14 ++++ .../templates/migrations-job.yaml | 2 +- .../tests/migrations-job_tests.yaml | 65 ++++++++++++++++++- deploy/charts/litellm-helm/values.yaml | 4 ++ 4 files changed, 83 insertions(+), 2 deletions(-) diff --git a/deploy/charts/litellm-helm/templates/_helpers.tpl b/deploy/charts/litellm-helm/templates/_helpers.tpl index a1eda28c679..25b02dd5f37 100644 --- a/deploy/charts/litellm-helm/templates/_helpers.tpl +++ b/deploy/charts/litellm-helm/templates/_helpers.tpl @@ -61,6 +61,20 @@ Create the name of the service account to use {{- end }} {{- end }} +{{/* +Create the service account name used by migration jobs. +When Helm hooks are enabled, pre-install/pre-upgrade hooks run before normal resources. +If this chart is creating the ServiceAccount, it is not yet available for the hook job, +so fall back to "default" (or an explicit override) to avoid a cyclic dependency. +*/}} +{{- define "litellm.migrationServiceAccountName" -}} +{{- if and .Values.migrationJob.hooks.helm.enabled .Values.serviceAccount.create }} +{{- default "default" .Values.migrationJob.serviceAccountName }} +{{- else }} +{{- include "litellm.serviceAccountName" . }} +{{- end }} +{{- end }} + {{/* Get redis service name */}} diff --git a/deploy/charts/litellm-helm/templates/migrations-job.yaml b/deploy/charts/litellm-helm/templates/migrations-job.yaml index 3459fa12d1c..8b93a60c1a3 100644 --- a/deploy/charts/litellm-helm/templates/migrations-job.yaml +++ b/deploy/charts/litellm-helm/templates/migrations-job.yaml @@ -34,7 +34,7 @@ spec: imagePullSecrets: {{- toYaml . | nindent 8 }} {{- end }} - serviceAccountName: {{ include "litellm.serviceAccountName" . }} + serviceAccountName: {{ include "litellm.migrationServiceAccountName" . }} {{- with .Values.migrationJob.extraInitContainers }} initContainers: {{- toYaml . | nindent 8 }} diff --git a/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml b/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml index 3a7bfa5eb0c..ee684c3c3d7 100644 --- a/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml +++ b/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml @@ -124,4 +124,67 @@ tests: - notContains: path: spec.template.spec.containers[0].env content: - name: DATABASE_URL \ No newline at end of file + name: DATABASE_URL + + - it: should use default service account for helm hooks when serviceAccount.create is true + template: migrations-job.yaml + set: + migrationJob: + enabled: true + hooks: + helm: + enabled: true + serviceAccount: + create: true + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: default + + - it: should use migrationJob.serviceAccountName override for helm hooks when serviceAccount.create is true + template: migrations-job.yaml + set: + migrationJob: + enabled: true + serviceAccountName: migration-sa + hooks: + helm: + enabled: true + serviceAccount: + create: true + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: migration-sa + + - it: should use chart service account when helm hooks are disabled + template: migrations-job.yaml + set: + migrationJob: + enabled: true + hooks: + helm: + enabled: false + serviceAccount: + create: true + name: my-custom-sa + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: my-custom-sa + + - it: should use pre-existing service account when helm hooks are enabled but serviceAccount.create is false + template: migrations-job.yaml + set: + migrationJob: + enabled: true + hooks: + helm: + enabled: true + serviceAccount: + create: false + name: pre-existing-sa + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: pre-existing-sa diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index d62f5b29c2b..bab909c8954 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -299,6 +299,10 @@ migrationJob: retries: 3 # Number of retries for the Job in case of failure backoffLimit: 4 # Backoff limit for Job restarts disableSchemaUpdate: false # Skip schema migrations for specific environments. When True, the job will exit with code 0. + # Optional service account for the migration job. + # Only used when migrationJob.hooks.helm.enabled=true and serviceAccount.create=true. + # In that case, pre-install/pre-upgrade hooks run before normal resources, so this defaults to "default". + serviceAccountName: "" annotations: {} ttlSecondsAfterFinished: 120 resources: {} From ce54c39051548dbb374979688ff4a9c68c639f8a Mon Sep 17 00:00:00 2001 From: Aarish Alam Date: Wed, 4 Mar 2026 09:55:35 +0530 Subject: [PATCH 04/58] Bug Fix: auto-inject prompt caching support for Gemini models (#21881) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add explicit caching to litellm proxy for gemini models via injection * fix: add missing `supports_function_calling` for deepinfra models All 55 deepinfra models that had `supports_tool_choice: true` were missing the `supports_function_calling` flag, causing `litellm.supports_function_calling()` to incorrectly return False. Fixes #22619 Co-Authored-By: Claude Opus 4.6 * Managed batches - Address PR bot comments from #22464 * feat(togetherai): add support for TogetherAI Qwen3.5-397B-A17B model * Agent Tracing - support context_id based trace id propogation + nested llm calls (#22626) * style(ui/): distinguish agent calls from llm calls on ui * feat: initial grouping working * feat: set stable contextid for a2a calls - allows for easily passing to downstream llm/mcp calls * feat(a2a_endpoints.py): fix tracing to avoid recreating logging objects for the same call allows stable trace id usage * fix(guardrail_endpoints): handle string ui_type values in _build_field_dict _build_field_dict unconditionally called .value on ui_type, which crashes for guardrail configs that use plain strings (e.g. BlockCodeExecutionGuardrailConfigModel uses "multiselect" and "percentage"). Now checks with hasattr before calling .value. Co-Authored-By: Claude Opus 4.6 * fix: propagate trace/session id from headers in MCP server calls Cherry-picked mcp_server/server.py fixes from 6feb9bab: adds get_chain_id_from_headers to extract x-litellm-trace-id / x-litellm-session-id from raw headers, and uses it in call_tool and list_tools to keep spend logs and tracing consistent with A2A. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 * [Feat] UI - Add Open in New Tab on leftnav Bar (#22731) * Add minimal dev_config.yaml for proxy development Co-authored-by: Ishaan Jaff * feat(ui): wrap left nav items in tags for open-in-new-tab support Nav items are now rendered as elements with proper href attributes, enabling right-click → 'Open in new tab', Ctrl/Cmd+click, and middle-click to open any sidebar page in a new browser tab. Normal clicks continue to use SPA navigation (no full page reload). Applied to both leftnav.tsx (query-param routing) and Sidebar2.tsx (Next.js file-based routing). Co-authored-by: Ishaan Jaff --------- Co-authored-by: Cursor Agent Co-authored-by: Ishaan Jaff * [Feat] Add Tool Policies for AI Gateway (#22732) * fix: fix ui render * fix: fix minor bugs * refactor: use prisma functions instead of raw sql (safer) * fix(add-new-tiles-to-tool-policies): allow developer to see what's available * feat: ensure tool allowlist runs correctly for tool names + mcp's * refactor: more ui improvements * feat: working key tool blocking * feat(tools): show tool logs * refactor: backend code improvements * refactor: improve log viewer for tools * fix: address PR review feedback for tool access control - Add missing blocked_tools column to root schema.prisma (schema drift) - Invalidate ToolPolicyRegistry after policy mutations so changes take effect immediately - Remove dead code: unused get_effective_policies, get_tool_policies_cached, and helpers Co-Authored-By: Claude Opus 4.6 * fix: race condition in permission resolution and remove duplicate allowlist check - Use atomic update_many with object_permission_id=None to prevent concurrent requests from creating orphaned permission rows and losing tool blocks - Remove duplicate allowed_tools enforcement from guardrail (already enforced in auth layer via check_tools_allowlist) - Move inline uuid import to module level Co-Authored-By: Claude Opus 4.6 * update to account for userAgent * UI - Add ToolDetails * input/output policy * LiteLLM_PolicyAttachmentTable * LiteLLM_PolicyAttachmentTable * fix: add _enqueue_tool_registry_upsert * fix: tool mgmt endpoints * tool mgmt endpoints * Update tests/test_litellm/proxy/db/test_tool_registry_writer.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update tests/test_litellm/proxy/db/test_tool_registry_writer.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update tests/test_litellm/proxy/db/test_tool_registry_writer.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix: sync root schema.prisma and fix test_tool_registry_writer for input/output policy - Migrate root schema.prisma LiteLLM_ToolTable from call_policy to input_policy/output_policy, add missing user_agent and last_used_at columns (now consistent with litellm/proxy/schema.prisma and litellm-proxy-extras) - Fix SpendLogToolIndex comment across all three schema files - Fix all call_policy references in test_tool_registry_writer.py: swapped update_tool_policy arguments, wrong get_tools_by_names return type assertions, _mock_tool_row setting call_policy instead of input_policy Addresses Greptile review feedback on PR #22732. Made-with: Cursor --------- Co-authored-by: Krrish Dholakia Co-authored-by: Claude Opus 4.6 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * feat(proxy): add key_alias, key_hash, requested_model DD APM span tags (#22710) * feat(proxy): add key_alias, key_hash, requested_model tags to DD APM spans * refactor(proxy): consolidate DD APM tag helpers into DDSpanTagger class * refactor(proxy): move DDSpanTagger to its own file litellm/proxy/dd_span_tagger.py --------- Co-authored-by: liweiguang Co-authored-by: Claude Opus 4.6 Co-authored-by: Ephrim Stanley Co-authored-by: Varad Khonde Co-authored-by: Krish Dholakia Co-authored-by: Sameer Kankute Co-authored-by: Ishaan Jaff Co-authored-by: Cursor Agent Co-authored-by: Ishaan Jaff Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- AGENTS.md | 3 + CLAUDE.md | 4 + dev_config.yaml | 13 + docs/my-website/docs/response_api.md | 2 +- .../migration.sql | 2 + .../migration.sql | 11 + .../litellm_proxy_extras/schema.prisma | 46 +- litellm/a2a_protocol/main.py | 91 ++-- litellm/files/main.py | 2 +- .../litellm_core_utils/get_litellm_params.py | 8 +- litellm/litellm_core_utils/litellm_logging.py | 306 +++++------ .../chat/guardrail_translation/handler.py | 10 +- .../guardrail_translation/base_translation.py | 7 + .../chat/guardrail_translation/handler.py | 13 + .../guardrail_translation/handler.py | 39 +- litellm/llms/vertex_ai/batches/handler.py | 5 +- .../vertex_ai_context_caching.py | 31 ++ .../llms/vertex_ai/files/transformation.py | 5 +- ...odel_prices_and_context_window_backup.json | 430 +++++++++++++-- .../proxy/_experimental/mcp_server/server.py | 17 +- .../out/{404.html => 404/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{budgets.html => budgets/index.html} | 0 .../{caching.html => caching/index.html} | 0 .../index.html} | 0 .../{old-usage.html => old-usage/index.html} | 0 .../{prompts.html => prompts/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{login.html => login/index.html} | 0 .../out/{logs.html => logs/index.html} | 0 .../{callback.html => callback/index.html} | 0 .../{model-hub.html => model-hub/index.html} | 0 .../{model_hub.html => model_hub/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{policies.html => policies/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{ui-theme.html => ui-theme/index.html} | 0 .../out/{teams.html => teams/index.html} | 0 .../{test-key.html => test-key/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{usage.html => usage/index.html} | 0 .../out/{users.html => users/index.html} | 0 .../index.html} | 0 litellm/proxy/_new_secret_config.yaml | 30 +- litellm/proxy/_types.py | 9 +- .../proxy/agent_endpoints/a2a_endpoints.py | 49 +- litellm/proxy/auth/auth_checks.py | 55 +- litellm/proxy/common_request_processing.py | 29 +- litellm/proxy/db/db_spend_update_writer.py | 155 +++--- litellm/proxy/db/spend_log_tool_index.py | 147 +++++ litellm/proxy/db/tool_registry_writer.py | 409 +++++++++++--- litellm/proxy/dd_span_tagger.py | 60 +++ .../proxy/guardrails/guardrail_endpoints.py | 8 +- .../tool_policy/tool_policy_guardrail.py | 208 +++++--- .../proxy/guardrails/tool_name_extraction.py | 85 +++ litellm/proxy/litellm_pre_call_utils.py | 88 +-- .../tool_management_endpoints.py | 501 +++++++++++++++++- litellm/proxy/proxy_server.py | 31 +- litellm/proxy/schema.prisma | 45 +- .../spend_tracking/spend_tracking_utils.py | 83 ++- litellm/proxy/utils.py | 18 +- litellm/types/tool_management.py | 66 ++- litellm/utils.py | 61 ++- model_prices_and_context_window.json | 178 +++++-- schema.prisma | 45 +- scripts/test_tool_allowlist_script.py | 116 ++++ .../test_anthropic_cache_control_hook.py | 100 ++++ .../test_vertex_ai_context_caching.py | 119 +++++ .../test_file_retrieve_provider_routing.py | 127 +++++ .../test_vertex_ai_files_transformation.py | 48 +- .../proxy/db/test_tool_registry_writer.py | 272 ++++++---- .../test_tool_policy_guardrail.py | 85 +-- .../guardrails/test_guardrail_endpoints.py | 43 +- .../proxy/test_common_request_processing.py | 66 ++- .../proxy/test_litellm_pre_call_utils.py | 100 ++-- .../proxy/test_tools_allowlist_enforcement.py | 200 +++++++ tests/test_litellm/test_utils.py | 45 ++ .../app/(dashboard)/components/Sidebar2.tsx | 25 +- ui/litellm-dashboard/src/app/page.tsx | 4 +- .../src/components/ToolDetail.tsx | 445 ++++++++++++++++ .../src/components/ToolPolicies.tsx | 304 +++++++---- .../components/ToolPolicies/PolicySelect.tsx | 92 ++++ .../src/components/ToolPoliciesView.tsx | 44 ++ .../src/components/leftnav.tsx | 44 +- .../src/components/networking.tsx | 166 +++++- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 39 +- .../src/components/view_logs/TypeBadges.tsx | 21 +- .../src/components/view_logs/columns.tsx | 25 +- .../src/components/view_logs/constants.ts | 3 + .../src/components/view_logs/index.tsx | 9 +- ui/litellm-dashboard/tsconfig.json | 2 +- 101 files changed, 4873 insertions(+), 1076 deletions(-) create mode 100644 dev_config.yaml create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql rename litellm/proxy/_experimental/out/{404.html => 404/index.html} (100%) rename litellm/proxy/_experimental/out/{_not-found.html => _not-found/index.html} (100%) rename litellm/proxy/_experimental/out/{api-reference.html => api-reference/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{api-playground.html => api-playground/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{budgets.html => budgets/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{caching.html => caching/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{claude-code-plugins.html => claude-code-plugins/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{old-usage.html => old-usage/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{prompts.html => prompts/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{tag-management.html => tag-management/index.html} (100%) rename litellm/proxy/_experimental/out/{guardrails.html => guardrails/index.html} (100%) rename litellm/proxy/_experimental/out/{login.html => login/index.html} (100%) rename litellm/proxy/_experimental/out/{logs.html => logs/index.html} (100%) rename litellm/proxy/_experimental/out/mcp/oauth/{callback.html => callback/index.html} (100%) rename litellm/proxy/_experimental/out/{model-hub.html => model-hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub.html => model_hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%) rename litellm/proxy/_experimental/out/{models-and-endpoints.html => models-and-endpoints/index.html} (100%) rename litellm/proxy/_experimental/out/{onboarding.html => onboarding/index.html} (100%) rename litellm/proxy/_experimental/out/{organizations.html => organizations/index.html} (100%) rename litellm/proxy/_experimental/out/{playground.html => playground/index.html} (100%) rename litellm/proxy/_experimental/out/{policies.html => policies/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{admin-settings.html => admin-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{logging-and-alerts.html => logging-and-alerts/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{router-settings.html => router-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{ui-theme.html => ui-theme/index.html} (100%) rename litellm/proxy/_experimental/out/{teams.html => teams/index.html} (100%) rename litellm/proxy/_experimental/out/{test-key.html => test-key/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{mcp-servers.html => mcp-servers/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{vector-stores.html => vector-stores/index.html} (100%) rename litellm/proxy/_experimental/out/{usage.html => usage/index.html} (100%) rename litellm/proxy/_experimental/out/{users.html => users/index.html} (100%) rename litellm/proxy/_experimental/out/{virtual-keys.html => virtual-keys/index.html} (100%) create mode 100644 litellm/proxy/db/spend_log_tool_index.py create mode 100644 litellm/proxy/dd_span_tagger.py create mode 100644 litellm/proxy/guardrails/tool_name_extraction.py create mode 100644 scripts/test_tool_allowlist_script.py create mode 100644 tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py create mode 100644 tests/test_litellm/proxy/test_tools_allowlist_enforcement.py create mode 100644 ui/litellm-dashboard/src/components/ToolDetail.tsx create mode 100644 ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.tsx create mode 100644 ui/litellm-dashboard/src/components/ToolPoliciesView.tsx diff --git a/AGENTS.md b/AGENTS.md index d43f41dbe30..546f2997bf5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -109,6 +109,8 @@ Key files: - `litellm/proxy/auth/` - Authentication logic - `litellm/proxy/management_endpoints/` - Admin API endpoints +**Database (proxy)**: Use Prisma model methods (`prisma_client.db..upsert`, `.find_many`, `.find_unique`, etc.), not raw SQL (`execute_raw`/`query_raw`). See COMMON PITFALLS for details. + ## MCP (MODEL CONTEXT PROTOCOL) SUPPORT LiteLLM supports MCP for agent workflows: @@ -176,6 +178,7 @@ When opening issues or pull requests, follow these templates: 5. **Dependencies**: Keep dependencies minimal and well-justified 6. **UI/Backend Contract Mismatch**: When adding a new entity type to the UI, always check whether the backend endpoint accepts a single value or an array. Match the UI control accordingly (single-select vs. multi-select) to avoid silently dropping user selections 7. **Missing Tests for New Entity Types**: When adding a new entity type (e.g., in `EntityUsage`, `UsageViewSelect`), always add corresponding tests in the existing test files and update any icon/component mocks +8. **Raw SQL in proxy DB code**: Do not use `execute_raw` or `query_raw` for proxy database access. Use Prisma model methods (e.g. `prisma_client.db.litellm_tooltable.upsert()`, `.find_many()`, `.find_unique()`) so behavior stays consistent with the schema, the client stays mockable in tests, and you avoid the pitfalls of hand-written SQL (parameter ordering, type casting, schema drift) 8. **Do not hardcode model-specific flags**: Put model-specific capability flags in `model_prices_and_context_window.json` and read them via `get_model_info` (or existing helpers like `supports_reasoning`). This prevents users from needing to upgrade LiteLLM each time a new model supports a feature. diff --git a/CLAUDE.md b/CLAUDE.md index 3b597fb8a90..c1eb75d2515 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -107,6 +107,10 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components: - Migration files auto-generated with `prisma migrate dev` - Always test migrations against both PostgreSQL and SQLite +### Proxy database access +- **Do not write raw SQL** for proxy DB operations. Use Prisma model methods instead of `execute_raw` / `query_raw`. +- Use the generated client: `prisma_client.db.` (e.g. `litellm_tooltable`, `litellm_usertable`) with `.upsert()`, `.find_many()`, `.find_unique()`, `.update()`, `.update_many()` as appropriate. This avoids schema/client drift, keeps code testable with simple mocks, and matches patterns used in spend logs and other proxy code. + ### Enterprise Features - Enterprise-specific code in `enterprise/` directory - Optional features enabled via environment variables diff --git a/dev_config.yaml b/dev_config.yaml new file mode 100644 index 00000000000..64e3c14703e --- /dev/null +++ b/dev_config.yaml @@ -0,0 +1,13 @@ +model_list: + - model_name: fake-openai-endpoint + litellm_params: + model: openai/fake-model + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + +general_settings: + master_key: sk-1234 + +litellm_settings: + drop_params: True + telemetry: False diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index a7cf61ef16a..76899a17ccb 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -930,7 +930,7 @@ For Responses API with load balancing across deployments with **different API ke Notes: - User-key affinity is keyed on `metadata.user_api_key_hash` (the API key hash). The OpenAI `user` request parameter is an end-user identifier and is intentionally not used for deployment affinity. -- Session-ID affinity is keyed on `metadata.session_id`. For proxy requests, this can be passed via the `x-litellm-session-id` HTTP header. For Python SDK requests, you can pass it via `litellm_metadata={"session_id": "value"}` in request args. +- Session-ID affinity is keyed on `metadata.session_id`. For proxy requests, this can be passed via the `x-litellm-session-id` or `x-litellm-trace-id` HTTP header (they are interchangeable for call chaining). For Python SDK requests, you can pass it via `litellm_metadata={"session_id": "value"}` in request args. - `user_api_key_hash` is already SHA-256, and is used as-is (no double hashing). - Affinity is scoped by a stable model identifier (the model-map key, e.g. `model_map_information.model_map_key`) so model aliases map to the same stickiness bucket. - The mapping TTL is controlled by `deployment_affinity_ttl_seconds` (configured on Router init / proxy startup). diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql new file mode 100644 index 00000000000..cba06684193 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "blocked_tools" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql new file mode 100644 index 00000000000..e3199679ce2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql @@ -0,0 +1,11 @@ +-- CreateTable +CREATE TABLE "LiteLLM_SpendLogToolIndex" ( + "request_id" TEXT NOT NULL, + "tool_name" TEXT NOT NULL, + "start_time" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_SpendLogToolIndex_pkey" PRIMARY KEY ("request_id","tool_name") +); + +-- CreateIndex +CREATE INDEX "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx" ON "LiteLLM_SpendLogToolIndex"("tool_name", "start_time"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index e0b28a4e012..5abe7a0a2b1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -260,6 +260,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -928,6 +929,16 @@ model LiteLLM_SpendLogGuardrailIndex { @@index([policy_id, start_time]) } +// Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production +model LiteLLM_SpendLogToolIndex { + request_id String + tool_name String // matches LiteLLM_ToolTable.tool_name; join for input_policy/output_policy etc. + start_time DateTime + + @@id([request_id, tool_name]) + @@index([tool_name, start_time]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) @@ -1065,26 +1076,31 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } -// Global tool registry - auto-discovered from LLM responses; admins set call_policy here +// Global tool registry - auto-discovered from LLM responses; admins set input_policy/output_policy here model LiteLLM_ToolTable { - tool_id String @id @default(uuid()) - tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" - origin String? // MCP server name or "user_defined" - call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" - call_count Int @default(0) // cumulative number of times this tool was seen - assignments Json? @default("{}") - key_hash String? // hash of the virtual key that first called this tool - team_id String? // team that first called this tool - key_alias String? // human-readable alias of the virtual key - created_at DateTime @default(now()) - created_by String? - updated_at DateTime @default(now()) @updatedAt - updated_by String? + tool_id String @id @default(uuid()) + tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" + origin String? // MCP server name or "user_defined" + input_policy String @default("untrusted") // "trusted" | "untrusted" | "blocked" + output_policy String @default("untrusted") // "trusted" | "untrusted" + call_count Int @default(0) // cumulative number of times this tool was seen + assignments Json? @default("{}") + key_hash String? // hash of the virtual key that first called this tool + team_id String? // team that first called this tool + key_alias String? // human-readable alias of the virtual key + user_agent String? // user-agent of the first request that discovered this tool + last_used_at DateTime? // timestamp of the most recent call + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? - @@index([call_policy]) + @@index([input_policy]) + @@index([output_policy]) @@index([team_id]) } +// Per-(tool, team/key) policy overrides. When present, override replaces global tool policy for that scope. //Unified Access Groups table for storing unified access groups model LiteLLM_AccessGroupTable { access_group_id String @id @default(uuid()) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 642dfaf023c..401b602fef5 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -24,11 +24,7 @@ from litellm.utils import client if TYPE_CHECKING: from a2a.client import A2AClient as A2AClientType - from a2a.types import ( - AgentCard, - SendMessageRequest, - SendStreamingMessageRequest, - ) + from a2a.types import AgentCard, SendMessageRequest, SendStreamingMessageRequest # Runtime imports with availability check A2A_SDK_AVAILABLE = False @@ -124,13 +120,48 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: litellm_logging_obj.model = model litellm_logging_obj.custom_llm_provider = custom_llm_provider litellm_logging_obj.model_call_details["model"] = model - litellm_logging_obj.model_call_details[ - "custom_llm_provider" - ] = custom_llm_provider + litellm_logging_obj.model_call_details["custom_llm_provider"] = ( + custom_llm_provider + ) return agent_name +async def _send_message_via_completion_bridge( + request: "SendMessageRequest", + custom_llm_provider: str, + api_base: Optional[str], + litellm_params: Dict[str, Any], +) -> LiteLLMSendMessageResponse: + """ + Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore). + + Requires request; api_base is optional for providers that derive endpoint from model. + """ + verbose_logger.info( + f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}" + ) + + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + params = ( + request.params.model_dump(mode="json") + if hasattr(request.params, "model_dump") + else dict(request.params) + ) + + response_dict = await A2ACompletionBridgeHandler.handle_non_streaming( + request_id=str(request.id), + params=params, + litellm_params=litellm_params, + api_base=api_base, + ) + + return LiteLLMSendMessageResponse.from_dict(response_dict) + + @client async def asend_message( a2a_client: Optional["A2AClientType"] = None, @@ -193,39 +224,21 @@ async def asend_message( ``` """ litellm_params = litellm_params or {} + logging_obj = kwargs.get("litellm_logging_obj") + trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None custom_llm_provider = litellm_params.get("custom_llm_provider") # Route through completion bridge if custom_llm_provider is set if custom_llm_provider: if request is None: raise ValueError("request is required for completion bridge") - # api_base is optional for providers that derive endpoint from model (e.g., bedrock/agentcore) - - verbose_logger.info( - f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}" - ) - - from litellm.a2a_protocol.litellm_completion_bridge.handler import ( - A2ACompletionBridgeHandler, - ) - - # Extract params from request - params = ( - request.params.model_dump(mode="json") - if hasattr(request.params, "model_dump") - else dict(request.params) - ) - - response_dict = await A2ACompletionBridgeHandler.handle_non_streaming( - request_id=str(request.id), - params=params, - litellm_params=litellm_params, + return await _send_message_via_completion_bridge( + request=request, + custom_llm_provider=custom_llm_provider, api_base=api_base, + litellm_params=litellm_params, ) - # Convert to LiteLLMSendMessageResponse - return LiteLLMSendMessageResponse.from_dict(response_dict) - # Standard A2A client flow if request is None: raise ValueError("request is required") @@ -236,11 +249,13 @@ async def asend_message( raise ValueError( "Either a2a_client or api_base is required for standard A2A flow" ) - trace_id = str(uuid.uuid4()) + trace_id = trace_id or str(uuid.uuid4()) extra_headers = {"X-LiteLLM-Trace-Id": trace_id} if agent_id: extra_headers["X-LiteLLM-Agent-Id"] = agent_id - a2a_client = await create_a2a_client(base_url=api_base, extra_headers=extra_headers) + a2a_client = await create_a2a_client( + base_url=api_base, extra_headers=extra_headers + ) # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None @@ -255,6 +270,10 @@ async def asend_message( ) card_url = getattr(agent_card, "url", None) if agent_card else None + context_id = trace_id or str(uuid.uuid4()) + if request.params.message.context_id is None: + request.params.message.context_id = context_id + # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL a2a_response = None for _ in range(2): # max 2 attempts: original + 1 retry @@ -606,7 +625,9 @@ async def create_a2a_client( if extra_headers: httpx_client.headers.update(extra_headers) - verbose_proxy_logger.debug(f"A2A client created with extra_headers={extra_headers}") + verbose_proxy_logger.debug( + f"A2A client created with extra_headers={extra_headers}" + ) # Resolve agent card resolver = A2ACardResolver( diff --git a/litellm/files/main.py b/litellm/files/main.py index 66d3a97468d..f0a8112fbdf 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -336,7 +336,7 @@ async def afile_retrieve( @client def file_retrieve( file_id: str, - custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "manus"] = "openai", + custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 36a8dfdb5a6..c91e4b6de1d 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -1,6 +1,5 @@ from typing import Optional - # Pre-define optional kwargs keys as frozenset for O(1) lookups # These are extracted from kwargs only if present, avoiding unnecessary .get() calls _OPTIONAL_KWARGS_KEYS = frozenset({ @@ -95,6 +94,13 @@ def get_litellm_params( litellm_request_debug: Optional[bool] = None, **kwargs, ) -> dict: + # Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining) + _meta = metadata or {} + if litellm_session_id is None: + litellm_session_id = _meta.get("session_id") or _meta.get("trace_id") + if litellm_trace_id is None: + litellm_trace_id = _meta.get("trace_id") or _meta.get("session_id") + # Build base dict with explicit parameters (always included) litellm_params = { "acompletion": acompletion, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5e5a6cea1b2..6f587abcdf1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -133,8 +133,8 @@ from ..integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger from ..integrations.azure_storage.azure_storage import AzureBlobStorageLogger from ..integrations.custom_prompt_management import CustomPromptManagement from ..integrations.datadog.datadog import DataDogLogger -from ..integrations.datadog.datadog_metrics import DatadogMetricsLogger from ..integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger +from ..integrations.datadog.datadog_metrics import DatadogMetricsLogger from ..integrations.dotprompt import DotpromptManager from ..integrations.dynamodb import DyanmoDBLogger from ..integrations.galileo import GalileoObserve @@ -352,9 +352,9 @@ class Logging(LiteLLMLoggingBaseClass): ) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[ - Any - ] = [] # for generating complete stream response + self.sync_streaming_chunks: List[Any] = ( + [] + ) # for generating complete stream response self.log_raw_request_response = log_raw_request_response # Initialize dynamic callbacks @@ -746,9 +746,9 @@ class Logging(LiteLLMLoggingBaseClass): prompt_spec=prompt_spec, dynamic_callback_params=dynamic_callback_params, ): - self.model_call_details[ - "prompt_integration" - ] = logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + logger.__class__.__name__ + ) return logger except Exception: # If check fails, continue to next logger @@ -816,9 +816,9 @@ class Logging(LiteLLMLoggingBaseClass): if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( non_default_params ): - self.model_call_details[ - "prompt_integration" - ] = anthropic_cache_control_logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + anthropic_cache_control_logger.__class__.__name__ + ) return anthropic_cache_control_logger ######################################################### @@ -830,9 +830,9 @@ class Logging(LiteLLMLoggingBaseClass): internal_usage_cache=None, llm_router=None, ) - self.model_call_details[ - "prompt_integration" - ] = vector_store_custom_logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + vector_store_custom_logger.__class__.__name__ + ) # Add to global callbacks so post-call hooks are invoked if ( vector_store_custom_logger @@ -892,9 +892,9 @@ class Logging(LiteLLMLoggingBaseClass): model ): # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model - self.model_call_details["litellm_params"][ - "api_base" - ] = self._get_masked_api_base(additional_args.get("api_base", "")) + self.model_call_details["litellm_params"]["api_base"] = ( + self._get_masked_api_base(additional_args.get("api_base", "")) + ) def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 # Log the exact input to the LLM API @@ -923,10 +923,10 @@ class Logging(LiteLLMLoggingBaseClass): try: # [Non-blocking Extra Debug Information in metadata] if turn_off_message_logging is True: - _metadata[ - "raw_request" - ] = "redacted by litellm. \ + _metadata["raw_request"] = ( + "redacted by litellm. \ 'litellm.turn_off_message_logging=True'" + ) else: curl_command = self._get_request_curl_command( api_base=additional_args.get("api_base", ""), @@ -937,34 +937,34 @@ class Logging(LiteLLMLoggingBaseClass): _metadata["raw_request"] = str(curl_command) # split up, so it's easier to parse in the UI - self.model_call_details[ - "raw_request_typed_dict" - ] = RawRequestTypedDict( - raw_request_api_base=str( - additional_args.get("api_base") or "" - ), - raw_request_body=self._get_raw_request_body( - additional_args.get("complete_input_dict", {}) - ), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, + self.model_call_details["raw_request_typed_dict"] = ( + RawRequestTypedDict( + raw_request_api_base=str( + additional_args.get("api_base") or "" + ), + raw_request_body=self._get_raw_request_body( + additional_args.get("complete_input_dict", {}) + ), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, + ) ) except Exception as e: - self.model_call_details[ - "raw_request_typed_dict" - ] = RawRequestTypedDict( - error=str(e), + self.model_call_details["raw_request_typed_dict"] = ( + RawRequestTypedDict( + error=str(e), + ) ) - _metadata[ - "raw_request" - ] = "Unable to Log \ + _metadata["raw_request"] = ( + "Unable to Log \ raw request: {}".format( - str(e) + str(e) + ) ) if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: @@ -1265,13 +1265,13 @@ class Logging(LiteLLMLoggingBaseClass): for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[ - MCPPostCallResponseObject - ] = await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, + response: Optional[MCPPostCallResponseObject] = ( + await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=post_mcp_tool_call_response_obj, + start_time=start_time, + end_time=end_time, + ) ) ###################################################################### # if any of the callbacks modify the response, use the modified response @@ -1466,9 +1466,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details[ - "response_cost_failure_debug_information" - ] = debug_info + self.model_call_details["response_cost_failure_debug_information"] = ( + debug_info + ) return None try: @@ -1494,9 +1494,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details[ - "response_cost_failure_debug_information" - ] = debug_info + self.model_call_details["response_cost_failure_debug_information"] = ( + debug_info + ) return None @@ -1652,10 +1652,8 @@ class Logging(LiteLLMLoggingBaseClass): result=logging_result ) - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - logging_result, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload(logging_result, start_time, end_time) ) if ( @@ -1734,9 +1732,9 @@ class Logging(LiteLLMLoggingBaseClass): end_time = datetime.datetime.now() if self.completion_start_time is None: self.completion_start_time = end_time - self.model_call_details[ - "completion_start_time" - ] = self.completion_start_time + self.model_call_details["completion_start_time"] = ( + self.completion_start_time + ) self.model_call_details["log_event_type"] = "successful_api_call" self.model_call_details["end_time"] = end_time @@ -1773,10 +1771,10 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, ) elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - result, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + result, start_time, end_time + ) ) if ( standard_logging_payload := self.model_call_details.get( @@ -1785,9 +1783,9 @@ class Logging(LiteLLMLoggingBaseClass): ) is not None: emit_standard_logging_payload(standard_logging_payload) elif standard_logging_object is not None: - self.model_call_details[ - "standard_logging_object" - ] = standard_logging_object + self.model_call_details["standard_logging_object"] = ( + standard_logging_object + ) else: self.model_call_details["response_cost"] = None @@ -1945,17 +1943,17 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( "Logging Details LiteLLM-Success Call streaming complete" ) - self.model_call_details[ - "complete_streaming_response" - ] = complete_streaming_response - self.model_call_details[ - "response_cost" - ] = self._response_cost_calculator(result=complete_streaming_response) + self.model_call_details["complete_streaming_response"] = ( + complete_streaming_response + ) + self.model_call_details["response_cost"] = ( + self._response_cost_calculator(result=complete_streaming_response) + ) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) ) if ( standard_logging_payload := self.model_call_details.get( @@ -2289,10 +2287,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} + self.model_call_details["complete_response"] = ( + self.model_call_details.get( + "complete_streaming_response", {} + ) ) result = self.model_call_details["complete_response"] openMeterLogger.log_success_event( @@ -2316,10 +2314,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} + self.model_call_details["complete_response"] = ( + self.model_call_details.get( + "complete_streaming_response", {} + ) ) result = self.model_call_details["complete_response"] @@ -2458,9 +2456,9 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") - self.model_call_details[ - "async_complete_streaming_response" - ] = complete_streaming_response + self.model_call_details["async_complete_streaming_response"] = ( + complete_streaming_response + ) try: if self.model_call_details.get("cache_hit", False) is True: @@ -2471,10 +2469,10 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details=self.model_call_details ) # base_model defaults to None if not set on model_info - self.model_call_details[ - "response_cost" - ] = self._response_cost_calculator( - result=complete_streaming_response + self.model_call_details["response_cost"] = ( + self._response_cost_calculator( + result=complete_streaming_response + ) ) verbose_logger.debug( @@ -2487,10 +2485,10 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = None ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) ) # print standard logging payload @@ -2517,10 +2515,8 @@ class Logging(LiteLLMLoggingBaseClass): # _success_handler_helper_fn if self.model_call_details.get("standard_logging_object") is None: ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - result, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload(result, start_time, end_time) ) # print standard logging payload @@ -2764,18 +2760,18 @@ class Logging(LiteLLMLoggingBaseClass): ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj={}, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="failure", - error_str=str(exception), - original_exception=exception, - standard_built_in_tools_params=self.standard_built_in_tools_params, + self.model_call_details["standard_logging_object"] = ( + get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj={}, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="failure", + error_str=str(exception), + original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) ) return start_time, end_time @@ -3739,9 +3735,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 service_name=arize_config.project_name, ) - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" + ) for callback in _in_memory_loggers: if ( isinstance(callback, ArizeLogger) @@ -3767,13 +3763,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" + ) else: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"openinference.project.name={arize_phoenix_config.project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"openinference.project.name={arize_phoenix_config.project_name}" + ) # Set Phoenix project name from environment variable phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) @@ -3781,19 +3777,19 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"{existing_attrs},openinference.project.name={phoenix_project_name}" + ) else: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"openinference.project.name={phoenix_project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"openinference.project.name={phoenix_project_name}" + ) # auth can be disabled on local deployments of arize phoenix if arize_phoenix_config.otlp_auth_headers is not None: - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = arize_phoenix_config.otlp_auth_headers + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + arize_phoenix_config.otlp_auth_headers + ) for callback in _in_memory_loggers: if ( @@ -3969,9 +3965,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 exporter="otlp_http", endpoint="https://langtrace.ai/api/trace", ) - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + f"api_key={os.getenv('LANGTRACE_API_KEY')}" + ) for callback in _in_memory_loggers: if ( isinstance(callback, OpenTelemetry) @@ -4204,8 +4200,7 @@ def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: litellm.logging_callback_manager.add_litellm_callback(phoenix_logger) verbose_logger.info( - "Auto-initialized Arize Phoenix logger alongside otel " - "(endpoint=%s)", + "Auto-initialized Arize Phoenix logger alongside otel " "(endpoint=%s)", arize_phoenix_config.endpoint, ) except Exception as e: @@ -4768,9 +4763,11 @@ class StandardLoggingPayloadSetup: ).model_dump() if isinstance(_raw, dict): if ResponseAPILoggingUtils._is_response_api_usage(_raw): - return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - _raw - ).model_dump() + return ( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + _raw + ).model_dump() + ) return _raw if isinstance(_raw, Usage): return _raw.model_dump() @@ -4884,10 +4881,10 @@ class StandardLoggingPayloadSetup: for key in StandardLoggingHiddenParams.__annotations__.keys(): if key in hidden_params: if key == "additional_headers": - clean_hidden_params[ - "additional_headers" - ] = StandardLoggingPayloadSetup.get_additional_headers( - hidden_params[key] + clean_hidden_params["additional_headers"] = ( + StandardLoggingPayloadSetup.get_additional_headers( + hidden_params[key] + ) ) else: clean_hidden_params[key] = hidden_params[key] # type: ignore @@ -5039,14 +5036,22 @@ class StandardLoggingPayloadSetup: dynamic_litellm_session_id = litellm_params.get("litellm_session_id") dynamic_litellm_trace_id = litellm_params.get("litellm_trace_id") + # Note: we recommend using `litellm_session_id` for session tracking # `litellm_trace_id` is an internal litellm param if dynamic_litellm_session_id: return str(dynamic_litellm_session_id) elif dynamic_litellm_trace_id: return str(dynamic_litellm_trace_id) - else: - return logging_obj.litellm_trace_id + # Fallback: use metadata.session_id or metadata.trace_id for call chaining + metadata = litellm_params.get("metadata") or {} + metadata_session_id = metadata.get("session_id") + metadata_trace_id = metadata.get("trace_id") + if metadata_session_id: + return str(metadata_session_id) + if metadata_trace_id: + return str(metadata_trace_id) + return logging_obj.litellm_trace_id @staticmethod def _get_user_agent_tags(proxy_server_request: dict) -> Optional[List[str]]: @@ -5502,9 +5507,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): ): for k, v in metadata["user_api_key_metadata"].items(): if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[ - k - ] = "scrubbed_by_litellm_for_sensitive_keys" + cleaned_user_api_key_metadata[k] = ( + "scrubbed_by_litellm_for_sensitive_keys" + ) else: cleaned_user_api_key_metadata[k] = v @@ -5616,4 +5621,3 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: model_parameters={"stream": True}, hidden_params=hidden_params, ) - diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 98650a238e9..a6df346e8a8 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -75,7 +75,7 @@ class AnthropicMessagesHandler(BaseTranslation): if messages is None: return data - chat_completion_compatible_request, tool_name_mapping = ( + chat_completion_compatible_request, _tool_name_mapping = ( LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( # Use a shallow copy to avoid mutating request data (pop on litellm_metadata). anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) @@ -141,6 +141,14 @@ class AnthropicMessagesHandler(BaseTranslation): return data + def extract_request_tool_names(self, data: dict) -> List[str]: + """Extract tool names from Anthropic messages request (tools[].name).""" + names: List[str] = [] + for tool in data.get("tools") or []: + if isinstance(tool, dict) and tool.get("name"): + names.append(str(tool["name"])) + return names + def _extract_input_text_and_images( self, message: Dict[str, Any], diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 7106c207bd6..a7982cb606e 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -98,3 +98,10 @@ class BaseTranslation(ABC): Optional to override in subclasses. """ return responses_so_far + + def extract_request_tool_names(self, data: dict) -> List[str]: + """ + Extract tool names from the request body for allowlist/policy checks. + Override in tool-capable handlers; default returns []. + """ + return [] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 67e9e42bc30..10b0b58b6ac 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -135,6 +135,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return data + def extract_request_tool_names(self, data: dict) -> List[str]: + """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" + names: List[str] = [] + for tool in data.get("tools") or []: + if isinstance(tool, dict) and tool.get("type") == "function": + fn = tool.get("function") + if isinstance(fn, dict) and fn.get("name"): + names.append(str(fn["name"])) + for fn in data.get("functions") or []: + if isinstance(fn, dict) and fn.get("name"): + names.append(str(fn["name"])) + return names + def _extract_inputs( self, message: Dict[str, Any], diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 6b092911d3c..7c3354cf88e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -30,27 +30,22 @@ Output: response.output is List[GenericResponseOutputItem] where each has: from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast -from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall +from openai.types.responses.response_function_tool_call import \ + ResponseFunctionToolCall from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, - OpenAiResponsesToChatCompletionStreamIterator, -) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation -from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, -) -from litellm.types.llms.openai import ( - ChatCompletionToolCallChunk, - ChatCompletionToolParam, -) -from litellm.types.responses.main import ( - GenericResponseOutputItem, - OutputFunctionToolCall, - OutputText, -) + OpenAiResponsesToChatCompletionStreamIterator) +from litellm.llms.base_llm.guardrail_translation.base_translation import \ + BaseTranslation +from litellm.responses.litellm_completion_transformation.transformation import \ + LiteLLMCompletionResponsesConfig +from litellm.types.llms.openai import (ChatCompletionToolCallChunk, + ChatCompletionToolParam) +from litellm.types.responses.main import (GenericResponseOutputItem, + OutputFunctionToolCall, OutputText) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -188,6 +183,18 @@ class OpenAIResponsesHandler(BaseTranslation): return data + def extract_request_tool_names(self, data: dict) -> List[str]: + """Extract tool names from Responses API request (tools[].name for function, tools[].server_label for mcp).""" + names: List[str] = [] + for tool in data.get("tools") or []: + if not isinstance(tool, dict): + continue + if tool.get("type") == "function" and tool.get("name"): + names.append(str(tool["name"])) + elif tool.get("type") == "mcp" and tool.get("server_label"): + names.append(str(tool["server_label"])) + return names + def _extract_and_transform_tools( self, tools: List[Dict[str, Any]], diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index ba3b5fb7a2c..5f1fefca963 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -115,9 +115,10 @@ class VertexAIBatchPrediction(VertexLLM): data=json.dumps(vertex_batch_request), ) except httpx.HTTPStatusError as e: - error_body = e.response.text if hasattr(e, 'response') else "N/A" + error_body = e.response.text litellm.verbose_logger.error( - f"Vertex AI batch create failed: status={e.response.status_code}, body={error_body[:1000]}" + "Vertex AI batch create failed: status=%s, body=%s", + e.response.status_code, error_body[:1000], ) raise if response.status_code != 200: diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index ed4d2d6a740..4450ae58349 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -4,13 +4,16 @@ import httpx import litellm from litellm.caching.caching import Cache, LiteLLMCacheType +from litellm.constants import MINIMUM_PROMPT_CACHE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, get_async_httpx_client, ) +from litellm._logging import verbose_logger from litellm.llms.openai.openai import AllMessageValues +from litellm.utils import is_prompt_caching_valid_prompt from litellm.types.llms.vertex_ai import ( CachedContentListAllResponseBody, VertexAICachedContentResponseObject, @@ -314,6 +317,20 @@ class ContextCachingEndpoints(VertexBase): if len(cached_messages) == 0: return messages, optional_params, None + # Gemini requires a minimum of 1024 tokens for context caching. + # Skip caching if the cached content is too small to avoid API errors. + if not is_prompt_caching_valid_prompt( + model=model, + messages=cached_messages, + custom_llm_provider=custom_llm_provider, + ): + verbose_logger.debug( + "Vertex AI context caching: cached content is below minimum token " + "count (%d). Skipping context caching.", + MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + ) + return messages, optional_params, None + tools = optional_params.pop("tools", None) ## AUTHORIZATION ## @@ -446,6 +463,20 @@ class ContextCachingEndpoints(VertexBase): if len(cached_messages) == 0: return messages, optional_params, None + # Gemini requires a minimum of 1024 tokens for context caching. + # Skip caching if the cached content is too small to avoid API errors. + if not is_prompt_caching_valid_prompt( + model=model, + messages=cached_messages, + custom_llm_provider=custom_llm_provider, + ): + verbose_logger.debug( + "Vertex AI context caching: cached content is below minimum token " + "count (%d). Skipping context caching.", + MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + ) + return messages, optional_params, None + tools = optional_params.pop("tools", None) ## AUTHORIZATION ## diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index f0493cd6be9..bf3ed5e6ac9 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -408,10 +408,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): file_id = "deleted" if hasattr(raw_response, "request") and raw_response.request: url = str(raw_response.request.url) - if "/o/" in url: + if "/b/" in url and "/o/" in url: import urllib.parse + bucket_part = url.split("/b/")[-1].split("/o/")[0] encoded_name = url.split("/o/")[-1].split("?")[0] - file_id = f"gs://{urllib.parse.unquote(encoded_name)}" + file_id = f"gs://{bucket_part}/{urllib.parse.unquote(encoded_name)}" return FileDeleted(id=file_id, deleted=True, object="file") def transform_list_files_request( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5de764c5cec..b92e2727979 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9779,6 +9779,122 @@ } ] }, + "dashscope/qwen3-max-2026-01-23": { + "litellm_provider": "dashscope", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "dashscope/qwen3-next-80b-a3b-instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "dashscope/qwen3-next-80b-a3b-thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "dashscope/qwen3-vl-235b-a22b-instruct": { + "input_cost_per_token": 4e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "dashscope/qwen3-vl-235b-a22b-thinking": { + "input_cost_per_token": 4e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "dashscope/qwen3-vl-32b-instruct": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.4e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "dashscope/qwen3-vl-32b-thinking": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "dashscope", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.87e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "dashscope/qwen3-vl-plus": { "litellm_provider": "dashscope", "max_input_tokens": 260096, @@ -10844,7 +10960,8 @@ "output_cost_per_token": 9e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": { "max_tokens": 131072, @@ -10854,7 +10971,8 @@ "output_cost_per_token": 1e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": { "max_tokens": 131072, @@ -10874,7 +10992,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen2.5-72B-Instruct": { "max_tokens": 32768, @@ -10884,7 +11003,8 @@ "output_cost_per_token": 3.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen2.5-7B-Instruct": { "max_tokens": 32768, @@ -10905,7 +11025,8 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-14B": { "max_tokens": 40960, @@ -10915,7 +11036,8 @@ "output_cost_per_token": 2.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B": { "max_tokens": 40960, @@ -10925,7 +11047,8 @@ "output_cost_per_token": 5.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": { "max_tokens": 262144, @@ -10935,7 +11058,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": { "max_tokens": 262144, @@ -10945,7 +11069,8 @@ "output_cost_per_token": 2.9e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-30B-A3B": { "max_tokens": 40960, @@ -10955,7 +11080,8 @@ "output_cost_per_token": 2.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-32B": { "max_tokens": 40960, @@ -10965,7 +11091,8 @@ "output_cost_per_token": 2.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": { "max_tokens": 262144, @@ -10975,7 +11102,8 @@ "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": { "max_tokens": 262144, @@ -10985,7 +11113,8 @@ "output_cost_per_token": 1.2e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Next-80B-A3B-Instruct": { "max_tokens": 262144, @@ -10995,7 +11124,8 @@ "output_cost_per_token": 1.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Next-80B-A3B-Thinking": { "max_tokens": 262144, @@ -11005,7 +11135,8 @@ "output_cost_per_token": 1.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": { "max_tokens": 8192, @@ -11056,7 +11187,8 @@ "cache_read_input_token_cost": 3.3e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/anthropic/claude-4-opus": { "max_tokens": 200000, @@ -11066,7 +11198,8 @@ "output_cost_per_token": 8.25e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/anthropic/claude-4-sonnet": { "max_tokens": 200000, @@ -11076,7 +11209,8 @@ "output_cost_per_token": 1.65e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1": { "max_tokens": 163840, @@ -11086,7 +11220,8 @@ "output_cost_per_token": 2.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528": { "max_tokens": 163840, @@ -11097,7 +11232,8 @@ "cache_read_input_token_cost": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": { "max_tokens": 32768, @@ -11107,7 +11243,8 @@ "output_cost_per_token": 3e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { "max_tokens": 131072, @@ -11127,7 +11264,8 @@ "output_cost_per_token": 2.7e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": { "max_tokens": 40960, @@ -11137,7 +11275,8 @@ "output_cost_per_token": 3e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3": { "max_tokens": 163840, @@ -11147,7 +11286,8 @@ "output_cost_per_token": 8.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3-0324": { "max_tokens": 163840, @@ -11157,7 +11297,8 @@ "output_cost_per_token": 8.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1": { "max_tokens": 163840, @@ -11169,7 +11310,8 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_reasoning": true + "supports_reasoning": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1-Terminus": { "max_tokens": 163840, @@ -11180,7 +11322,8 @@ "cache_read_input_token_cost": 2.16e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.0-flash-001": { "deprecation_date": "2026-06-01", @@ -11191,7 +11334,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.5-flash": { "max_tokens": 1000000, @@ -11201,7 +11345,8 @@ "output_cost_per_token": 2.5e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -11211,7 +11356,8 @@ "output_cost_per_token": 1e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -11221,7 +11367,8 @@ "output_cost_per_token": 1e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-27b-it": { "max_tokens": 131072, @@ -11231,7 +11378,8 @@ "output_cost_per_token": 1.6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-4b-it": { "max_tokens": 131072, @@ -11241,7 +11389,8 @@ "output_cost_per_token": 8e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": { "max_tokens": 131072, @@ -11261,7 +11410,8 @@ "output_cost_per_token": 2e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct": { "max_tokens": 131072, @@ -11271,7 +11421,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": { "max_tokens": 131072, @@ -11281,6 +11432,7 @@ "output_cost_per_token": 3.9e-07, "litellm_provider": "deepinfra", "mode": "chat", + "supports_function_calling": true, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { @@ -11291,7 +11443,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": { "max_tokens": 327680, @@ -11301,7 +11454,8 @@ "output_cost_per_token": 3e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-Guard-3-8B": { "max_tokens": 131072, @@ -11331,7 +11485,8 @@ "output_cost_per_token": 6e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": { "max_tokens": 131072, @@ -11341,7 +11496,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { "max_tokens": 131072, @@ -11351,7 +11507,8 @@ "output_cost_per_token": 2.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": { "max_tokens": 131072, @@ -11361,7 +11518,8 @@ "output_cost_per_token": 5e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { "max_tokens": 131072, @@ -11371,7 +11529,8 @@ "output_cost_per_token": 3e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/microsoft/WizardLM-2-8x22B": { "max_tokens": 65536, @@ -11391,7 +11550,8 @@ "output_cost_per_token": 1.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": { "max_tokens": 131072, @@ -11401,7 +11561,8 @@ "output_cost_per_token": 4e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": { "max_tokens": 32768, @@ -11411,7 +11572,8 @@ "output_cost_per_token": 8e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": { "max_tokens": 128000, @@ -11421,7 +11583,8 @@ "output_cost_per_token": 2e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { "max_tokens": 32768, @@ -11431,7 +11594,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/moonshotai/Kimi-K2-Instruct": { "max_tokens": 131072, @@ -11441,7 +11605,8 @@ "output_cost_per_token": 2e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/moonshotai/Kimi-K2-Instruct-0905": { "max_tokens": 262144, @@ -11452,7 +11617,8 @@ "cache_read_input_token_cost": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": { "max_tokens": 131072, @@ -11462,7 +11628,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/Llama-3.3-Nemotron-Super-49B-v1.5": { "max_tokens": 131072, @@ -11472,7 +11639,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { "max_tokens": 131072, @@ -11482,7 +11650,8 @@ "output_cost_per_token": 1.6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/openai/gpt-oss-120b": { "max_tokens": 131072, @@ -11492,7 +11661,8 @@ "output_cost_per_token": 4.5e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/openai/gpt-oss-20b": { "max_tokens": 131072, @@ -11502,7 +11672,8 @@ "output_cost_per_token": 1.5e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/zai-org/GLM-4.5": { "max_tokens": 131072, @@ -11512,7 +11683,8 @@ "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepseek/deepseek-chat": { "cache_creation_input_token_cost": 0.0, @@ -25806,6 +25978,30 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "openrouter/anthropic/claude-sonnet-4.6": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "source": "https://openrouter.ai/anthropic/claude-sonnet-4.6", + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 159 + }, "openrouter/anthropic/claude-opus-4.5": { "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, @@ -26156,6 +26352,39 @@ "supports_web_search": true, "tpm": 800000 }, + "openrouter/google/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "source": "https://openrouter.ai/google/gemini-3.1-pro-preview", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 1.875e-06, "litellm_provider": "openrouter", @@ -26533,6 +26762,29 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/openai/gpt-5.1-codex-max": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://openrouter.ai/openai/gpt-5.1-codex-max", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "openrouter/openai/gpt-5.2": { "input_cost_per_image": 0, "cache_read_input_token_cost": 1.75e-07, @@ -26687,6 +26939,19 @@ "supports_tool_choice": true, "supports_function_calling": true }, + "openrouter/qwen/qwen3-coder-plus": { + "input_cost_per_token": 1e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://openrouter.ai/qwen/qwen3-coder-plus", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/qwen/qwen3-235b-a22b-2507": { "input_cost_per_token": 7.1e-08, "litellm_provider": "openrouter", @@ -26822,6 +27087,19 @@ "supports_vision": true, "supports_prompt_caching": false }, + "openrouter/z-ai/glm-5": { + "input_cost_per_token": 8e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 202752, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.56e-06, + "source": "https://openrouter.ai/z-ai/glm-5", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/minimax/minimax-m2.1": { "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1.2e-06, @@ -29736,6 +30014,18 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", @@ -34315,6 +34605,36 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "zai/glm-5": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, + "zai/glm-5-code": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 5e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.7": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 1.1e-07, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5b3d5bd60e2..cdb26acb658 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -5,7 +5,6 @@ LiteLLM MCP Server Routes import asyncio import contextlib - import traceback import uuid from datetime import datetime @@ -44,7 +43,10 @@ from litellm.proxy._experimental.mcp_server.utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.ip_address_utils import IPAddressUtils -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + get_chain_id_from_headers, +) from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall @@ -331,6 +333,11 @@ if MCP_AVAILABLE: try: # Create a body date for logging body_data = {"name": name, "arguments": arguments} + # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A) + chain_id = get_chain_id_from_headers(raw_headers) + if chain_id: + body_data["litellm_trace_id"] = chain_id + body_data["litellm_session_id"] = chain_id request = Request( scope={ @@ -884,6 +891,10 @@ if MCP_AVAILABLE: # This is intentionally minimal: only async_success_handler / post_call_failure_hook rules_obj = Rules() list_tools_call_id = str(uuid.uuid4()) + # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) + effective_litellm_trace_id = litellm_trace_id or get_chain_id_from_headers( + raw_headers + ) spend_logs_metadata: Dict[str, Any] = { "mcp_operation": "list_tools", } @@ -896,7 +907,7 @@ if MCP_AVAILABLE: "model": "MCP: list_tools", "call_type": CallTypes.list_mcp_tools.value, "litellm_call_id": list_tools_call_id, - "litellm_trace_id": litellm_trace_id, + "litellm_trace_id": effective_litellm_trace_id, "metadata": { "spend_logs_metadata": spend_logs_metadata, }, diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/_not-found.html b/litellm/proxy/_experimental/out/_not-found/index.html similarity index 100% rename from litellm/proxy/_experimental/out/_not-found.html rename to litellm/proxy/_experimental/out/_not-found/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 6b84d90a327..508c1c94659 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -23,33 +23,11 @@ model_list: guardrails: - - guardrail_name: "airline-competitor-intent" - guardrail_id: "airline-competitor-intent" + - guardrail_name: "tool_policy" litellm_params: - guardrail: litellm_content_filter - mode: pre_call - default_on: false - competitor_intent_config: - brand_self: - - emirates - - ek - competitors: - - qatar airways - - qatar - - etihad - locations: - - qatar - - doha - - doh - competitor_aliases: - qatar airways: [qr, doha airline] - qatar: [qr] - policy: - competitor_comparison: refuse - possible_competitor_comparison: reframe - threshold_high: 0.70 - threshold_medium: 0.45 - threshold_low: 0.30 + guardrail: tool_policy + mode: [pre_call, post_call] + default_on: true mcp_servers: my_http_server: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8d49020461d..42b48446e7a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -77,6 +77,7 @@ class SupportedDBObjectType(str, enum.Enum): PASS_THROUGH_ENDPOINTS = "pass_through_endpoints" PROMPTS = "prompts" MODEL_COST_MAP = "model_cost_map" + TOOLS = "tools" def __str__(self): return str(self.value) @@ -2133,7 +2134,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): user_header_mappings: Optional[List[UserHeaderMapping]] = None supported_db_objects: Optional[List[SupportedDBObjectType]] = Field( None, - description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map'. If not set, all objects are loaded (default behavior).", + description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'tools'. If not set, all objects are loaded (default behavior).", ) user_mcp_management_mode: Optional[UserMCPManagementMode] = Field( None, @@ -3377,6 +3378,11 @@ class ProxyErrorTypes(str, enum.Enum): Team member is already in team """ + tool_access_denied = "tool_access_denied" + """ + Tool is not in the allowed tools list for this key/team + """ + @classmethod def get_model_access_error_type_for_object( cls, object_type: Literal["key", "user", "team", "org", "project"] @@ -4161,6 +4167,7 @@ class ToolDiscoveryQueueItem(TypedDict, total=False): key_hash: Optional[str] # hash of virtual key that triggered discovery team_id: Optional[str] # team that triggered discovery key_alias: Optional[str] # human-readable key alias + user_agent: Optional[str] # HTTP User-Agent of the caller class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 7f30277ebca..6bcee14f29e 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -69,6 +69,7 @@ async def _handle_stream_message( from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE if not A2A_SDK_AVAILABLE: + async def _error_stream(): yield json.dumps( { @@ -106,7 +107,12 @@ async def _handle_stream_message( proxy_server_request=proxy_server_request, ) - if use_proxy_hooks and user_api_key_dict is not None and request_data is not None and proxy_logging_obj is not None: + if ( + use_proxy_hooks + and user_api_key_dict is not None + and request_data is not None + and proxy_logging_obj is not None + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -119,20 +125,27 @@ async def _handle_stream_message( return json.dumps(obj) + "\n" def _ndjson_error(proxy_exc: Any) -> str: - return json.dumps( - { - "jsonrpc": "2.0", - "id": request_id, - "error": { - "code": -32603, - "message": getattr( - proxy_exc, "message", f"Streaming error: {proxy_exc!s}" - ), - }, - } - ) + "\n" + return ( + json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, + "message": getattr( + proxy_exc, + "message", + f"Streaming error: {proxy_exc!s}", + ), + }, + } + ) + + "\n" + ) - async for line in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + async for ( + line + ) in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=a2a_stream, user_api_key_dict=user_api_key_dict, request_data=request_data, @@ -151,7 +164,12 @@ async def _handle_stream_message( yield json.dumps(chunk) + "\n" except Exception as e: verbose_proxy_logger.exception(f"Error streaming A2A response: {e}") - if use_proxy_hooks and proxy_logging_obj is not None and user_api_key_dict is not None and request_data is not None: + if ( + use_proxy_hooks + and proxy_logging_obj is not None + and user_api_key_dict is not None + and request_data is not None + ): transformed_exception = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -382,6 +400,7 @@ async def invoke_agent_a2a( agent_id=agent.agent_id, metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), + litellm_logging_obj=logging_obj, ) response = await proxy_logging_obj.post_call_success_hook( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 91ac58215ab..79e5f78f68e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -58,6 +58,10 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.guardrails.tool_name_extraction import ( + TOOL_CAPABLE_CALL_TYPES, + extract_request_tool_names, +) from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router @@ -220,7 +224,48 @@ async def _run_project_checks( ) -async def common_checks( +async def check_tools_allowlist( + request_body: dict, + valid_token: Optional[UserAPIKeyAuth], + team_object: Optional[LiteLLM_TeamTable], + route: str, +) -> None: + """ + Enforce key/team tool allowlist (metadata.allowed_tools). No DB in hot path — + effective allowlist is read from valid_token.metadata and valid_token.team_metadata. + Raises ProxyException with tool_access_denied if a tool is not allowed. + """ + from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, + ) + + if valid_token is None: + return + call_types = get_call_types_for_route(route) + if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types): + return + tool_names = extract_request_tool_names(route, request_body) + if not tool_names: + return + key_meta = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} + team_meta = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {} + key_allowed = key_meta.get("allowed_tools") + team_allowed = team_meta.get("allowed_tools") + effective = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed + if not isinstance(effective, list) or len(effective) == 0: + return + allowed_set = {str(t) for t in effective} + disallowed = [n for n in tool_names if n not in allowed_set] + if disallowed: + raise ProxyException( + message=f"Tool(s) {disallowed} are not in the allowed tools list for this key/team.", + type=ProxyErrorTypes.tool_access_denied, + param="tools", + code=status.HTTP_403_FORBIDDEN, + ) + + +async def common_checks( # noqa: PLR0915 request_body: dict, team_object: Optional[LiteLLM_TeamTable], user_object: Optional[LiteLLM_UserTable], @@ -477,6 +522,14 @@ async def common_checks( valid_token=valid_token, ) + # 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path) + await check_tools_allowlist( + request_body=request_body, + valid_token=valid_token, + team_object=team_object, + route=route, + ) + return True diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1269f58213a..e91f3fcd270 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -29,7 +29,7 @@ from litellm.constants import ( MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, STREAM_SSE_DATA_PREFIX, ) -from litellm.litellm_core_utils.dd_tracing import set_active_span_tag, tracer +from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, @@ -41,6 +41,7 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, ) +from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging from litellm.router import Router @@ -245,26 +246,6 @@ async def create_response( ) -def _add_dd_apm_tags_for_litellm_call_id(litellm_call_id: Optional[str]) -> None: - """ - Attach LiteLLM call id to the active Datadog APM span. - - This enables searching APM traces by LiteLLM call id returned in - `x-litellm-call-id`. - """ - if not litellm_call_id: - return - - try: - set_active_span_tag("litellm.call_id", str(litellm_call_id)) - except Exception: - # Tagging is best-effort and should never impact request processing. - verbose_proxy_logger.debug( - "Failed to tag active ddtrace span with litellm.call_id", - exc_info=True, - ) - - def _override_openai_response_model( *, response_obj: Any, @@ -662,7 +643,11 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_call_id"] = request.headers.get( "x-litellm-call-id", str(uuid.uuid4()) ) - _add_dd_apm_tags_for_litellm_call_id(self.data.get("litellm_call_id")) + DDSpanTagger.tag_call_id(self.data.get("litellm_call_id")) + DDSpanTagger.tag_request( + user_api_key_dict=user_api_key_dict, + requested_model=self.data.get("model"), + ) ### AUTO STREAM USAGE TRACKING ### # If always_include_stream_usage is enabled and this is a streaming request diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 0c25424ceaa..4c96e079c9e 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -13,49 +13,36 @@ import random import time import traceback from datetime import datetime, timedelta, timezone -from typing import ( - TYPE_CHECKING, - Any, - Dict, - List, - Literal, - Optional, - Union, - cast, - overload, -) +from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, + cast, overload) import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache, RedisCache from litellm.constants import DB_SPEND_UPDATE_JOB_NAME from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, - BaseDailySpendTransaction, - DailyAgentSpendTransaction, - DailyEndUserSpendTransaction, - DailyOrganizationSpendTransaction, - DailyTagSpendTransaction, - DailyTeamSpendTransaction, - DailyUserSpendTransaction, - DBSpendUpdateTransactions, - Litellm_EntityType, - LiteLLM_UserTable, - SpendLogsMetadata, - SpendLogsPayload, - SpendUpdateQueueItem, - ToolDiscoveryQueueItem, -) -from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( - DailySpendUpdateQueue, -) -from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager -from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer -from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue -from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( - ToolDiscoveryQueue, -) +from litellm.proxy._types import (DB_CONNECTION_ERROR_TYPES, + BaseDailySpendTransaction, + DailyAgentSpendTransaction, + DailyEndUserSpendTransaction, + DailyOrganizationSpendTransaction, + DailyTagSpendTransaction, + DailyTeamSpendTransaction, + DailyUserSpendTransaction, + DBSpendUpdateTransactions, + Litellm_EntityType, LiteLLM_UserTable, + SpendLogsMetadata, SpendLogsPayload, + SpendUpdateQueueItem, ToolDiscoveryQueueItem) +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \ + DailySpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import \ + PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import \ + RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.spend_update_queue import \ + SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import \ + ToolDiscoveryQueue from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING if TYPE_CHECKING: @@ -104,12 +91,10 @@ class DBSpendUpdateWriter: end_time: Optional[datetime], response_cost: Optional[float], ): - from litellm.proxy.proxy_server import ( - disable_spend_logs, - litellm_proxy_budget_name, - prisma_client, - user_api_key_cache, - ) + from litellm.proxy.proxy_server import (disable_spend_logs, + litellm_proxy_budget_name, + prisma_client, + user_api_key_cache) from litellm.proxy.utils import ProxyUpdateSpend, hash_token try: @@ -124,9 +109,8 @@ class DBSpendUpdateWriter: hashed_token = token ## CREATE SPEND LOG PAYLOAD ## - from litellm.proxy.spend_tracking.spend_tracking_utils import ( - get_logging_payload, - ) + from litellm.proxy.spend_tracking.spend_tracking_utils import \ + get_logging_payload payload = get_logging_payload( kwargs=kwargs, @@ -230,6 +214,7 @@ class DBSpendUpdateWriter: _litellm_params = kwargs.get("litellm_params") or {} _metadata = _litellm_params.get("metadata") or {} key_alias = _metadata.get("user_api_key_alias") or None + user_agent = _metadata.get("user_agent") or None def _enqueue(tool_name: str, origin: str = "user_defined") -> None: self.tool_discovery_queue.add_update( @@ -239,17 +224,20 @@ class DBSpendUpdateWriter: key_hash=hashed_token, team_id=team_id, key_alias=key_alias, + user_agent=user_agent, ) ) # --- MCP tool calls --- sl_object = kwargs.get("standard_logging_object") if sl_object is not None: - mcp_metadata = ( - sl_object.get("metadata", {}) or {} - ).get("mcp_tool_call_metadata") + mcp_metadata = (sl_object.get("metadata", {}) or {}).get( + "mcp_tool_call_metadata" + ) if mcp_metadata and isinstance(mcp_metadata, dict): - tool_name = mcp_metadata.get("namespaced_tool_name") or mcp_metadata.get("name") + tool_name = mcp_metadata.get( + "namespaced_tool_name" + ) or mcp_metadata.get("name") mcp_server_name = mcp_metadata.get("mcp_server_name") if tool_name: _enqueue(tool_name, origin=mcp_server_name or "user_defined") @@ -280,7 +268,9 @@ class DBSpendUpdateWriter: _enqueue(name) # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) --- - if completion_response is not None and hasattr(completion_response, "choices"): + if completion_response is not None and hasattr( + completion_response, "choices" + ): for choice in completion_response.choices or []: message = getattr(choice, "message", None) if message is None: @@ -768,19 +758,46 @@ class DBSpendUpdateWriter: daily_end_user_spend_update_transactions, daily_agent_spend_update_transactions, daily_tag_spend_update_transactions, - ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + ) = ( + await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + ) if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d", - len(db_spend_update_transactions.get("key_list_transactions") or {}), - len(db_spend_update_transactions.get("user_list_transactions") or {}), - len(db_spend_update_transactions.get("team_list_transactions") or {}), - len(db_spend_update_transactions.get("org_list_transactions") or {}), - len(db_spend_update_transactions.get("end_user_list_transactions") or {}), - len(db_spend_update_transactions.get("team_member_list_transactions") or {}), - len(db_spend_update_transactions.get("tag_list_transactions") or {}), + len( + db_spend_update_transactions.get("key_list_transactions") + or {} + ), + len( + db_spend_update_transactions.get("user_list_transactions") + or {} + ), + len( + db_spend_update_transactions.get("team_list_transactions") + or {} + ), + len( + db_spend_update_transactions.get("org_list_transactions") + or {} + ), + len( + db_spend_update_transactions.get( + "end_user_list_transactions" + ) + or {} + ), + len( + db_spend_update_transactions.get( + "team_member_list_transactions" + ) + or {} + ), + len( + db_spend_update_transactions.get("tag_list_transactions") + or {} + ), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -985,10 +1002,8 @@ class DBSpendUpdateWriter: Commits all the spend `UPDATE` transactions to the Database """ - from litellm.proxy.utils import ( - ProxyUpdateSpend, - _raise_failed_update_spend_exception, - ) + from litellm.proxy.utils import (ProxyUpdateSpend, + _raise_failed_update_spend_exception) ### UPDATE USER TABLE ### user_list_transactions = db_spend_update_transactions["user_list_transactions"] @@ -1523,14 +1538,14 @@ class DBSpendUpdateWriter: # Add cache-related fields if they exist if "cache_read_input_tokens" in transaction: - common_data[ - "cache_read_input_tokens" - ] = transaction.get("cache_read_input_tokens", 0) + common_data["cache_read_input_tokens"] = ( + transaction.get("cache_read_input_tokens", 0) + ) if "cache_creation_input_tokens" in transaction: - common_data[ - "cache_creation_input_tokens" - ] = transaction.get( - "cache_creation_input_tokens", 0 + common_data["cache_creation_input_tokens"] = ( + transaction.get( + "cache_creation_input_tokens", 0 + ) ) if entity_type == "tag" and "request_id" in transaction: diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py new file mode 100644 index 00000000000..6e8c63675e6 --- /dev/null +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -0,0 +1,147 @@ +""" +Track tool usage for the dashboard: insert into SpendLogToolIndex when spend logs +are written, so "last N requests for tool X" and "how is this tool called in production" +queries are fast. +""" + +from datetime import datetime, timezone +from typing import Any, Dict, List, Set + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.proxy.utils import PrismaClient + + +def _add_tool_calls_to_set(tool_calls: Any, out: Set[str]) -> None: + """Extract tool names from OpenAI-style tool_calls list into out.""" + if not isinstance(tool_calls, list): + return + for tc in tool_calls: + if not isinstance(tc, dict): + continue + fn = tc.get("function") + if isinstance(fn, dict): + name = fn.get("name") + if name and isinstance(name, str) and name.strip(): + out.add(name.strip()) + + +def _parse_tool_names_from_payload(payload: Dict[str, Any]) -> Set[str]: + """ + Extract deduplicated tool names from a spend log payload. + Sources: mcp_namespaced_tool_name, response (tool_calls), proxy_server_request (tools). + """ + tool_names: Set[str] = set() + + # Top-level MCP tool name (single tool per request for that flow) + mcp_name = payload.get("mcp_namespaced_tool_name") + if mcp_name and isinstance(mcp_name, str) and mcp_name.strip(): + tool_names.add(mcp_name.strip()) + + # Response: OpenAI-style tool_calls[].function.name or choices[0].message.tool_calls + response_raw = payload.get("response") + if response_raw: + response_obj = ( + safe_json_loads(response_raw, default=None) + if isinstance(response_raw, str) + else response_raw + ) + if isinstance(response_obj, dict): + _add_tool_calls_to_set(response_obj.get("tool_calls"), tool_names) + choices = response_obj.get("choices") + if isinstance(choices, list) and choices: + msg = choices[0].get("message") if isinstance(choices[0], dict) else None + if isinstance(msg, dict): + _add_tool_calls_to_set(msg.get("tool_calls"), tool_names) + + # Request body: tools[].function.name + request_raw = payload.get("proxy_server_request") + if request_raw: + request_obj = ( + safe_json_loads(request_raw, default=None) + if isinstance(request_raw, str) + else request_raw + ) + if isinstance(request_obj, dict): + body = request_obj.get("body", request_obj) + if isinstance(body, dict): + request_obj = body + if isinstance(request_obj, dict): + tools = request_obj.get("tools") + if isinstance(tools, list): + for t in tools: + if isinstance(t, dict): + fn = t.get("function") + if isinstance(fn, dict): + name = fn.get("name") + if name and isinstance(name, str) and name.strip(): + tool_names.add(name.strip()) + + return tool_names + + +async def process_spend_logs_tool_usage( + prisma_client: PrismaClient, + logs_to_process: List[Dict[str, Any]], +) -> None: + """ + After spend logs are written: insert SpendLogToolIndex rows from each payload. + Extracts tool names from mcp_namespaced_tool_name, response tool_calls, and + proxy_server_request tools. + """ + if not logs_to_process: + return + + index_rows: List[Dict[str, Any]] = [] + + for payload in logs_to_process: + request_id = payload.get("request_id") + start_time = payload.get("startTime") + if not request_id or not start_time: + continue + if isinstance(start_time, str): + try: + start_time = datetime.fromisoformat( + start_time.replace("Z", "+00:00") + ) + except (ValueError, TypeError): + continue + if start_time.tzinfo is None: + start_time = start_time.replace(tzinfo=timezone.utc) + + tool_names = _parse_tool_names_from_payload(payload) + for tool_name in tool_names: + index_rows.append({ + "request_id": request_id, + "tool_name": tool_name, + "start_time": start_time, + }) + + if not index_rows: + return + + try: + index_data = [] + for r in index_rows: + st = r["start_time"] + if isinstance(st, str): + try: + st = datetime.fromisoformat(st.replace("Z", "+00:00")) + except (ValueError, TypeError): + continue + if st.tzinfo is None: + st = st.replace(tzinfo=timezone.utc) + index_data.append({ + "request_id": r["request_id"], + "tool_name": r["tool_name"], + "start_time": st, + }) + if index_data: + await prisma_client.db.litellm_spendlogtoolindex.create_many( + data=index_data, + skip_duplicates=True, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Tool usage tracking (SpendLogToolIndex) failed (non-fatal): %s", e + ) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 4e0a8095a08..0eda012d515 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -2,36 +2,64 @@ DB helpers for LiteLLM_ToolTable — the global tool registry. Tools are auto-discovered from LLM responses and upserted here. -Admins use the management endpoints to read and update call_policy. - -NOTE: Uses raw SQL (query_raw / execute_raw) instead of Prisma model methods -because the generated Prisma Python client may not have LiteLLM_ToolTable -when running against an older generated schema. +Admins use the management endpoints to read and update input_policy / output_policy. """ import uuid from datetime import datetime, timezone -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem -from litellm.types.tool_management import LiteLLM_ToolTableRow, ToolCallPolicy +from litellm.types.tool_management import ( + LiteLLM_ToolTableRow, + ToolPolicyOverrideRow, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient -def _row_to_model(row: dict) -> LiteLLM_ToolTableRow: +def _row_to_model(row: Union[dict, Any]) -> LiteLLM_ToolTableRow: + """Convert a Prisma model instance or dict to LiteLLM_ToolTableRow.""" + model_dump = getattr(row, "model_dump", None) + if callable(model_dump): + row = model_dump() + elif not isinstance(row, dict): + row = { + k: getattr(row, k, None) + for k in ( + "tool_id", + "tool_name", + "origin", + "input_policy", + "output_policy", + "call_count", + "assignments", + "key_hash", + "team_id", + "key_alias", + "user_agent", + "last_used_at", + "created_at", + "updated_at", + "created_by", + "updated_by", + ) + } return LiteLLM_ToolTableRow( tool_id=row.get("tool_id", ""), tool_name=row.get("tool_name", ""), origin=row.get("origin"), - call_policy=row.get("call_policy", "untrusted"), + input_policy=row.get("input_policy") or "untrusted", + output_policy=row.get("output_policy") or "untrusted", call_count=int(row.get("call_count") or 0), assignments=row.get("assignments"), key_hash=row.get("key_hash"), team_id=row.get("team_id"), key_alias=row.get("key_alias"), + user_agent=row.get("user_agent"), + last_used_at=row.get("last_used_at"), created_at=row.get("created_at"), updated_at=row.get("updated_at"), created_by=row.get("created_by"), @@ -44,10 +72,10 @@ async def batch_upsert_tools( items: List[ToolDiscoveryQueueItem], ) -> None: """ - Batch-upsert tool registry rows via raw SQL. + Batch-upsert tool registry rows via Prisma. - On first insert: sets call_policy = "untrusted" (schema default), call_count = 1. - On conflict: increments call_count; preserves existing call_policy. + On first insert: sets input_policy/output_policy = "untrusted" (default), call_count = 1. + On conflict: increments call_count; preserves existing policies. """ if not items: return @@ -55,6 +83,8 @@ async def batch_upsert_tools( data = [item for item in items if item.get("tool_name")] if not data: return + now = datetime.now(timezone.utc) + table = prisma_client.db.litellm_tooltable for item in data: tool_name = item.get("tool_name", "") origin = item.get("origin") or "user_defined" @@ -62,49 +92,52 @@ async def batch_upsert_tools( key_hash = item.get("key_hash") team_id = item.get("team_id") key_alias = item.get("key_alias") - now = datetime.now(timezone.utc).isoformat() - await prisma_client.db.execute_raw( - 'INSERT INTO "LiteLLM_ToolTable" ' - "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias, created_at, updated_at) " - "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8, $8) " - "ON CONFLICT (tool_name) DO UPDATE SET " - "call_count = \"LiteLLM_ToolTable\".call_count + 1, " - "updated_at = $8", - tool_name, - origin, - created_by, - key_hash, - team_id, - key_alias, - str(uuid.uuid4()), - now, + user_agent = item.get("user_agent") + await table.upsert( + where={"tool_name": tool_name}, + data={ + "create": { + "tool_id": str(uuid.uuid4()), + "tool_name": tool_name, + "origin": origin, + "input_policy": "untrusted", + "output_policy": "untrusted", + "call_count": 1, + "created_by": created_by, + "updated_by": created_by, + "key_hash": key_hash, + "team_id": team_id, + "key_alias": key_alias, + "user_agent": user_agent, + "last_used_at": now, + }, + "update": { + "call_count": {"increment": 1}, + "updated_at": now, + "last_used_at": now, + }, + }, ) verbose_proxy_logger.debug( "tool_registry_writer: upserted %d tool(s)", len(data) ) except Exception as e: - verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e) + verbose_proxy_logger.error( + "tool_registry_writer batch_upsert_tools error: %s", e + ) async def list_tools( prisma_client: "PrismaClient", - call_policy: Optional[ToolCallPolicy] = None, + input_policy: Optional[str] = None, ) -> List[LiteLLM_ToolTableRow]: - """Return all tools, optionally filtered by call_policy.""" + """Return all tools, optionally filtered by input_policy.""" try: - if call_policy is not None: - rows = await prisma_client.db.query_raw( - 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' - 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' - 'FROM "LiteLLM_ToolTable" WHERE call_policy = $1 ORDER BY created_at DESC', - call_policy, - ) - else: - rows = await prisma_client.db.query_raw( - 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' - 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' - 'FROM "LiteLLM_ToolTable" ORDER BY created_at DESC', - ) + where = {"input_policy": input_policy} if input_policy is not None else {} + rows = await prisma_client.db.litellm_tooltable.find_many( + where=where, + order={"created_at": "desc"}, + ) return [_row_to_model(row) for row in rows] except Exception as e: verbose_proxy_logger.error("tool_registry_writer list_tools error: %s", e) @@ -117,15 +150,12 @@ async def get_tool( ) -> Optional[LiteLLM_ToolTableRow]: """Return a single tool row by tool_name.""" try: - rows = await prisma_client.db.query_raw( - 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' - 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' - 'FROM "LiteLLM_ToolTable" WHERE tool_name = $1', - tool_name, + row = await prisma_client.db.litellm_tooltable.find_unique( + where={"tool_name": tool_name}, ) - if not rows: + if row is None: return None - return _row_to_model(rows[0]) + return _row_to_model(row) except Exception as e: verbose_proxy_logger.error("tool_registry_writer get_tool error: %s", e) return None @@ -134,46 +164,279 @@ async def get_tool( async def update_tool_policy( prisma_client: "PrismaClient", tool_name: str, - call_policy: ToolCallPolicy, updated_by: Optional[str], + input_policy: Optional[str] = None, + output_policy: Optional[str] = None, ) -> Optional[LiteLLM_ToolTableRow]: - """Update the call_policy for a tool. Upserts the row if it does not exist yet.""" + """Update input_policy and/or output_policy for a tool. Upserts the row if it does not exist yet.""" try: _updated_by = updated_by or "system" - now = datetime.now(timezone.utc).isoformat() - await prisma_client.db.execute_raw( - 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by, created_at, updated_at) ' - "VALUES ($4, $1, $2, $3, $3, $5, $5) " - "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5", - tool_name, - call_policy, - _updated_by, - str(uuid.uuid4()), - now, + now = datetime.now(timezone.utc) + + create_data: dict = { + "tool_id": str(uuid.uuid4()), + "tool_name": tool_name, + "input_policy": input_policy or "untrusted", + "output_policy": output_policy or "untrusted", + "created_by": _updated_by, + "updated_by": _updated_by, + "created_at": now, + "updated_at": now, + } + update_data: dict = { + "updated_by": _updated_by, + "updated_at": now, + } + if input_policy is not None: + update_data["input_policy"] = input_policy + if output_policy is not None: + update_data["output_policy"] = output_policy + + await prisma_client.db.litellm_tooltable.upsert( + where={"tool_name": tool_name}, + data={ + "create": create_data, + "update": update_data, + }, ) return await get_tool(prisma_client, tool_name) except Exception as e: - verbose_proxy_logger.error("tool_registry_writer update_tool_policy error: %s", e) + verbose_proxy_logger.error( + "tool_registry_writer update_tool_policy error: %s", e + ) return None async def get_tools_by_names( prisma_client: "PrismaClient", tool_names: List[str], -) -> Dict[str, str]: +) -> Dict[str, Tuple[str, str]]: """ - Return a {tool_name: call_policy} map for the given tool names. - Used by the policy enforcement guardrail — single batch query, never N+1. + Return a {tool_name: (input_policy, output_policy)} map for the given tool names. """ if not tool_names: return {} try: - placeholders = ", ".join(f"${i+1}" for i in range(len(tool_names))) - rows = await prisma_client.db.query_raw( - f'SELECT tool_name, call_policy FROM "LiteLLM_ToolTable" WHERE tool_name IN ({placeholders})', - *tool_names, + rows = await prisma_client.db.litellm_tooltable.find_many( + where={"tool_name": {"in": tool_names}}, ) - return {row["tool_name"]: row["call_policy"] for row in rows} + return { + row.tool_name: ( + getattr(row, "input_policy", "untrusted") or "untrusted", + getattr(row, "output_policy", "untrusted") or "untrusted", + ) + for row in rows + } except Exception as e: - verbose_proxy_logger.error("tool_registry_writer get_tools_by_names error: %s", e) + verbose_proxy_logger.error( + "tool_registry_writer get_tools_by_names error: %s", e + ) return {} + + +async def list_overrides_for_tool( + prisma_client: "PrismaClient", + tool_name: str, +) -> List[ToolPolicyOverrideRow]: + """ + Return override-like rows for a tool by finding object permissions that have + this tool in blocked_tools, then resolving each permission to key/team scope for display. + """ + out: List[ToolPolicyOverrideRow] = [] + try: + perms = await prisma_client.db.litellm_objectpermissiontable.find_many( + where={"blocked_tools": {"has": tool_name}}, + include={ + "verification_tokens": True, + "teams": True, + }, + ) + for perm in perms: + op_id = getattr(perm, "object_permission_id", None) or "" + tokens = getattr(perm, "verification_tokens", []) or [] + teams = getattr(perm, "teams", []) or [] + for t in tokens: + out.append( + ToolPolicyOverrideRow( + override_id=op_id, + tool_name=tool_name, + team_id=None, + key_hash=getattr(t, "token", None), + input_policy="blocked", + key_alias=getattr(t, "key_alias", None), + created_at=None, + updated_at=None, + ) + ) + for team in teams: + out.append( + ToolPolicyOverrideRow( + override_id=op_id, + tool_name=tool_name, + team_id=getattr(team, "team_id", None), + key_hash=None, + input_policy="blocked", + key_alias=getattr(team, "team_alias", None), + created_at=None, + updated_at=None, + ) + ) + return out + except Exception as e: + verbose_proxy_logger.error( + "tool_registry_writer list_overrides_for_tool error: %s", e + ) + return [] + + +class ToolPolicyRegistry: + """ + In-memory registry of tool policies synced from DB. + Hot path uses get_effective_policies only — no DB, no cache. + """ + + def __init__(self) -> None: + self._tool_input_policies: Dict[str, str] = {} + self._tool_output_policies: Dict[str, str] = {} + self._blocked_tools_by_op_id: Dict[str, List[str]] = {} + self._initialized: bool = False + + def is_initialized(self) -> bool: + return self._initialized + + async def sync_tool_policy_from_db(self, prisma_client: "PrismaClient") -> None: + """Load all tool policies and object-permission blocked_tools from DB.""" + try: + tools = await prisma_client.db.litellm_tooltable.find_many() + self._tool_input_policies = { + row.tool_name: getattr(row, "input_policy", "untrusted") or "untrusted" + for row in tools + } + self._tool_output_policies = { + row.tool_name: getattr(row, "output_policy", "untrusted") or "untrusted" + for row in tools + } + + perms = await prisma_client.db.litellm_objectpermissiontable.find_many() + self._blocked_tools_by_op_id = {} + for row in perms: + op_id = getattr(row, "object_permission_id", None) + blocked = getattr(row, "blocked_tools", None) or [] + if op_id: + self._blocked_tools_by_op_id[op_id] = list(blocked) + + self._initialized = True + verbose_proxy_logger.info( + "ToolPolicyRegistry: synced %d tool policies and %d object permissions from DB", + len(self._tool_input_policies), + len(self._blocked_tools_by_op_id), + ) + except Exception as e: + verbose_proxy_logger.exception( + "ToolPolicyRegistry sync_tool_policy_from_db error: %s", e + ) + raise + + def get_input_policy(self, tool_name: str) -> str: + return self._tool_input_policies.get(tool_name, "untrusted") + + def get_output_policy(self, tool_name: str) -> str: + return self._tool_output_policies.get(tool_name, "untrusted") + + def get_effective_policies( + self, + tool_names: List[str], + object_permission_id: Optional[str] = None, + team_object_permission_id: Optional[str] = None, + ) -> Dict[str, str]: + """ + Return effective input_policy per tool from in-memory state. + If tool is in key or team blocked_tools -> "blocked", else global input_policy or "untrusted". + """ + if not tool_names: + return {} + blocked: set = set() + for op_id in (object_permission_id, team_object_permission_id): + if op_id and op_id.strip(): + blocked.update( + self._blocked_tools_by_op_id.get(op_id.strip(), []) + ) + result: Dict[str, str] = {} + for name in tool_names: + if name in blocked: + result[name] = "blocked" + else: + result[name] = self._tool_input_policies.get(name, "untrusted") + return result + + +_tool_policy_registry: Optional[ToolPolicyRegistry] = None + + +def get_tool_policy_registry() -> ToolPolicyRegistry: + """Return the global ToolPolicyRegistry singleton.""" + global _tool_policy_registry + if _tool_policy_registry is None: + _tool_policy_registry = ToolPolicyRegistry() + return _tool_policy_registry + + +async def add_tool_to_object_permission_blocked( + prisma_client: "PrismaClient", + object_permission_id: str, + tool_name: str, +) -> bool: + """Add tool_name to the permission's blocked_tools if not already present.""" + if not object_permission_id or not tool_name: + return False + try: + row = await prisma_client.db.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": object_permission_id}, + ) + if row is None: + return False + current = list(getattr(row, "blocked_tools", []) or []) + if tool_name in current: + return True + current.append(tool_name) + await prisma_client.db.litellm_objectpermissiontable.update( + where={"object_permission_id": object_permission_id}, + data={"blocked_tools": current}, + ) + return True + except Exception as e: + verbose_proxy_logger.error( + "tool_registry_writer add_tool_to_object_permission_blocked error: %s", e + ) + return False + + +async def remove_tool_from_object_permission_blocked( + prisma_client: "PrismaClient", + object_permission_id: str, + tool_name: str, +) -> bool: + """Remove tool_name from the permission's blocked_tools. Returns False if tool was not in list.""" + if not object_permission_id or not tool_name: + return False + try: + row = await prisma_client.db.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": object_permission_id}, + ) + if row is None: + return False + current = list(getattr(row, "blocked_tools", []) or []) + if tool_name not in current: + return False + current = [t for t in current if t != tool_name] + await prisma_client.db.litellm_objectpermissiontable.update( + where={"object_permission_id": object_permission_id}, + data={"blocked_tools": current}, + ) + return True + except Exception as e: + verbose_proxy_logger.error( + "tool_registry_writer remove_tool_from_object_permission_blocked error: %s", + e, + ) + return False diff --git a/litellm/proxy/dd_span_tagger.py b/litellm/proxy/dd_span_tagger.py new file mode 100644 index 00000000000..08b7d928d0e --- /dev/null +++ b/litellm/proxy/dd_span_tagger.py @@ -0,0 +1,60 @@ +from typing import Optional + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.dd_tracing import set_active_span_tag +from litellm.proxy._types import UserAPIKeyAuth + + +class DDSpanTagger: + """Best-effort helpers for tagging the active Datadog APM span with LiteLLM request metadata.""" + + @staticmethod + def tag_call_id(litellm_call_id: Optional[str]) -> None: + """ + Attach LiteLLM call id to the active Datadog APM span. + + This enables searching APM traces by LiteLLM call id returned in + `x-litellm-call-id`. + """ + if not litellm_call_id: + return + try: + set_active_span_tag("litellm.call_id", str(litellm_call_id)) + except Exception: + verbose_proxy_logger.debug( + "Failed to tag active ddtrace span with litellm.call_id", + exc_info=True, + ) + + @staticmethod + def tag_request( + user_api_key_dict: UserAPIKeyAuth, + requested_model: Optional[str], + ) -> None: + """ + Attach key and model tags to the active Datadog APM span. + + Tags set (all best-effort, skipped when value is absent): + - ``litellm.key_alias`` — human-readable alias for the API key + - ``litellm.key_hash`` — hashed API key (safe to log; never the raw secret) + - ``litellm.requested_model``— model name as sent by the client + + Use cases: + - Trace all requests from a specific user/key: filter by ``litellm.key_alias`` or + ``litellm.key_hash``. + - Trace all requests for a specific model: filter by ``litellm.requested_model``. + + Note: key_alias / key_hash are not available for unauthenticated (e.g. 401) requests. + """ + try: + if user_api_key_dict.key_alias: + set_active_span_tag("litellm.key_alias", str(user_api_key_dict.key_alias)) + if user_api_key_dict.token: + set_active_span_tag("litellm.key_hash", str(user_api_key_dict.token)) + if requested_model: + set_active_span_tag("litellm.requested_model", str(requested_model)) + except Exception: + verbose_proxy_logger.debug( + "Failed to tag active ddtrace span with key/model tags", + exc_info=True, + ) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index c6a709534e1..4c866a24991 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1624,11 +1624,11 @@ def _build_field_dict( # Determine the field type from annotation field_type = _get_field_type_from_annotation(field_annotation) - # Check for custom UI type override (ui_type preferred; "type" leaks into OpenAPI and breaks schema) - field_json_schema_extra = getattr(field, "json_schema_extra", {}) or {} + # Check for custom UI type override + field_json_schema_extra = getattr(field, "json_schema_extra", {}) if field_json_schema_extra and "ui_type" in field_json_schema_extra: - ut = field_json_schema_extra["ui_type"] - field_type = ut if isinstance(ut, str) else getattr(ut, "value", ut) + ui_type = field_json_schema_extra["ui_type"] + field_type = ui_type.value if hasattr(ui_type, "value") else ui_type elif field_json_schema_extra and "type" in field_json_schema_extra: field_type = field_json_schema_extra["type"] diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index 87558566c42..368948414e9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -1,13 +1,16 @@ """ Tool Policy Guardrail -Reads call_policy from LiteLLM_ToolTable and enforces it on LLM requests/responses. +Reads input_policy / output_policy from LiteLLM_ToolTable and enforces them. -Policy values: - "trusted" - allow through (no action) - "untrusted" - allow through (no action; default for newly discovered tools) +Input policy values: + "untrusted" - allow through (default for newly discovered tools) + "trusted" - only allow if conversation contains no untrusted tool output "blocked" - raise HTTPException, preventing the tool call - "dual_llm" - (Phase 3) send to second LLM for verification; currently treated as allowed + +Output policy values: + "untrusted" - output may be tainted (default) + "trusted" - output is verified safe Configuration in proxy config YAML: guardrails: @@ -15,25 +18,18 @@ Configuration in proxy config YAML: litellm_params: guardrail: tool_policy mode: post_call - -or both pre and post call: - - guardrail_name: "tool_policy" - litellm_params: - guardrail: tool_policy - mode: during_call # runs before LLM and on response """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.caching.dual_cache import DualCache -from litellm.constants import TOOL_POLICY_CACHE_TTL_SECONDS from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.proxy.guardrails.tool_name_extraction import extract_request_tool_names from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -43,12 +39,71 @@ if TYPE_CHECKING: GUARDRAIL_NAME = "tool_policy" +def _get_request_object_permission_ids( + request_data: dict, +) -> Tuple[Optional[str], Optional[str]]: + """Extract object_permission_id and team_object_permission_id from request_data.""" + if not request_data: + return None, None + for key in ("litellm_metadata", "metadata"): + meta = request_data.get(key) + if not isinstance(meta, dict): + continue + auth = meta.get("user_api_key_auth") + if auth is not None and hasattr(auth, "object_permission_id"): + key_op = getattr(auth, "object_permission_id", None) + team_op = getattr(auth, "team_object_permission_id", None) + if key_op is not None or team_op is not None: + return ( + str(key_op).strip() if key_op else None, + str(team_op).strip() if team_op else None, + ) + key_op = meta.get("user_api_key_object_permission_id") + team_op = meta.get("user_api_key_team_object_permission_id") + if key_op is not None or team_op is not None: + return ( + str(key_op).strip() if key_op else None, + str(team_op).strip() if team_op else None, + ) + return None, None + + +def _get_request_route_from_data(request_data: dict) -> Optional[str]: + """Get request route from request_data (metadata or top-level).""" + route = request_data.get("user_api_key_request_route") + if route: + return route + meta = request_data.get("metadata") or request_data.get("litellm_metadata") or {} + return meta.get("user_api_key_request_route") + + +def _resolve_tool_names_from_messages(messages: List[dict]) -> Dict[str, str]: + """ + Build a map of tool_call_id -> tool_name from assistant messages' tool_calls. + Used to resolve which tool produced each tool result in the conversation. + """ + mapping: Dict[str, str] = {} + for msg in messages: + if msg.get("role") != "assistant": + continue + tool_calls = msg.get("tool_calls") or [] + for tc in tool_calls: + if isinstance(tc, dict): + tc_id = tc.get("id") + fn = (tc.get("function") or {}).get("name") + else: + tc_id = getattr(tc, "id", None) + fn_obj = getattr(tc, "function", None) + fn = getattr(fn_obj, "name", None) if fn_obj else None + if tc_id and fn: + mapping[tc_id] = fn + return mapping + + class ToolPolicyGuardrail(CustomGuardrail): """ - Guardrail that enforces per-tool call policies stored in LiteLLM_ToolTable. - - Tools with call_policy="blocked" are rejected before/after the LLM call. - Tools with call_policy="trusted" or "untrusted" pass through unchanged. + Guardrail that enforces per-tool input/output policies from the in-memory + ToolPolicyRegistry (synced from DB). """ def __init__(self, **kwargs: Any) -> None: @@ -59,7 +114,6 @@ class ToolPolicyGuardrail(CustomGuardrail): GuardrailEventHooks.during_call, ] super().__init__(**kwargs) - self._policy_cache: DualCache = DualCache() @log_guardrail_information async def apply_guardrail( @@ -70,12 +124,7 @@ class ToolPolicyGuardrail(CustomGuardrail): logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: """ - Enforce tool policies on both request tools and response tool_calls. - - - input_type="request": check inputs["tools"] (tool definitions in the LLM request) - - input_type="response": check inputs["tool_calls"] (tool_calls in the LLM response) - - Raises HTTPException (400) if any tool is "blocked". + Enforce input_policy and output_policy trust chain on request tools / response tool_calls. """ if input_type == "request": tools = inputs.get("tools") or [] @@ -86,7 +135,11 @@ class ToolPolicyGuardrail(CustomGuardrail): and isinstance(t.get("function"), dict) and t["function"].get("name") ] - else: # response + if not tool_names: + route = _get_request_route_from_data(request_data) + if route: + tool_names = extract_request_tool_names(route, request_data) + else: tool_calls = inputs.get("tool_calls") or [] tool_names = [] for tc in tool_calls: @@ -101,12 +154,25 @@ class ToolPolicyGuardrail(CustomGuardrail): if not tool_names: return inputs - policy_map = await self._get_policies_cached(tool_names) + object_permission_id, team_object_permission_id = ( + _get_request_object_permission_ids(request_data) + ) + from litellm.proxy.db.tool_registry_writer import get_tool_policy_registry + registry = get_tool_policy_registry() + if not registry.is_initialized(): + return inputs + + # Stage 1: Check for blocked tools (input_policy=blocked or per-key/team override) + policy_map = registry.get_effective_policies( + tool_names, + object_permission_id=object_permission_id, + team_object_permission_id=team_object_permission_id, + ) blocked = [name for name in tool_names if policy_map.get(name) == "blocked"] if blocked: verbose_proxy_logger.warning( - "ToolPolicyGuardrail: blocking tool(s) %s (policy=blocked)", blocked + "ToolPolicyGuardrail: blocking tool(s) %s (input_policy=blocked)", blocked ) raise HTTPException( status_code=400, @@ -117,47 +183,47 @@ class ToolPolicyGuardrail(CustomGuardrail): }, ) + # Stage 2: Trust chain enforcement (response path only) + # For each tool with input_policy=trusted, check if conversation + # contains output from tools with output_policy=untrusted + if input_type == "response": + trusted_input_tools = [ + name for name in tool_names if policy_map.get(name) == "trusted" + ] + if trusted_input_tools: + messages = request_data.get("messages") or [] + tc_id_to_name = _resolve_tool_names_from_messages(messages) + + untrusted_sources: List[str] = [] + for msg in messages: + if msg.get("role") != "tool": + continue + tool_call_id = msg.get("tool_call_id") + source_tool = tc_id_to_name.get(tool_call_id, "") if tool_call_id else "" + if not source_tool: + continue + if registry.get_output_policy(source_tool) == "untrusted": + if source_tool not in untrusted_sources: + untrusted_sources.append(source_tool) + + if untrusted_sources: + verbose_proxy_logger.warning( + "ToolPolicyGuardrail: trust chain violation — %s require trusted input " + "but conversation has untrusted output from %s", + trusted_input_tools, + untrusted_sources, + ) + raise HTTPException( + status_code=400, + detail={ + "error": "Violated tool policy", + "blocked_tools": trusted_input_tools, + "untrusted_sources": untrusted_sources, + "message": ( + f"{', '.join(trusted_input_tools)} requires trusted input but " + f"conversation contains untrusted output from {', '.join(untrusted_sources)}." + ), + }, + ) + return inputs - - async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]: - """ - Batch-fetch call_policy for the given tool names. - - Caches per individual tool name (not per combination) so that adding - a new tool to a request doesn't invalidate the cached policies for all - the other tools already in the cache. - """ - from litellm.proxy.db.tool_registry_writer import get_tools_by_names - from litellm.proxy.proxy_server import prisma_client - - if not tool_names or prisma_client is None: - return {} - - result: Dict[str, str] = {} - cache_misses: List[str] = [] - - for name in tool_names: - cached = await self._policy_cache.async_get_cache(f"tool_policy:{name}") - if cached is not None and isinstance(cached, str): - result[name] = cached - else: - cache_misses.append(name) - - if cache_misses: - fetched = await get_tools_by_names( - prisma_client=prisma_client, tool_names=cache_misses - ) - for name, policy in fetched.items(): - result[name] = policy - await self._policy_cache.async_set_cache( - key=f"tool_policy:{name}", - value=policy, - ttl=TOOL_POLICY_CACHE_TTL_SECONDS, - ) - verbose_proxy_logger.debug( - "ToolPolicyGuardrail: fetched %d policies from DB (cache hits: %d)", - len(cache_misses), - len(tool_names) - len(cache_misses), - ) - - return result diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py new file mode 100644 index 00000000000..db24fa2277c --- /dev/null +++ b/litellm/proxy/guardrails/tool_name_extraction.py @@ -0,0 +1,85 @@ +""" +Extract tool names from request body by route/call type. + +Used by auth (check_tools_allowlist) and ToolPolicyGuardrail so tool-format +knowledge lives in one place. Uses guardrail translation handlers where available, +with standalone extractors for generate_content and MCP. +""" + +from typing import Any, Dict, List + +from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route +from litellm.llms import load_guardrail_translation_mappings +from litellm.types.utils import CallTypes + +# Call types that have no guardrail translation handler; we use standalone extractors +STANDALONE_EXTRACTORS: Dict[str, Any] = {} + + +def _extract_generate_content_tool_names(data: dict) -> List[str]: + """Google generateContent: tools[].functionDeclarations[].name""" + names: List[str] = [] + for tool in data.get("tools") or []: + if not isinstance(tool, dict): + continue + for decl in tool.get("functionDeclarations") or []: + if isinstance(decl, dict) and decl.get("name"): + names.append(str(decl["name"])) + return names + + +def _extract_mcp_tool_names(data: dict) -> List[str]: + """MCP call_tool: name or mcp_tool_name in body""" + names: List[str] = [] + name = data.get("name") or data.get("mcp_tool_name") + if name: + names.append(str(name)) + return names + + +def _register_standalone_extractors() -> None: + if STANDALONE_EXTRACTORS: + return + STANDALONE_EXTRACTORS[CallTypes.generate_content.value] = _extract_generate_content_tool_names + STANDALONE_EXTRACTORS[CallTypes.agenerate_content.value] = _extract_generate_content_tool_names + STANDALONE_EXTRACTORS[CallTypes.call_mcp_tool.value] = _extract_mcp_tool_names + + +# Tool-capable call types (routes that can send tools in the request) +TOOL_CAPABLE_CALL_TYPES = frozenset({ + CallTypes.completion.value, + CallTypes.acompletion.value, + CallTypes.responses.value, + CallTypes.aresponses.value, + CallTypes.anthropic_messages.value, + CallTypes.generate_content.value, + CallTypes.agenerate_content.value, + CallTypes.call_mcp_tool.value, +}) + + +def extract_request_tool_names(route: str, data: dict) -> List[str]: + """ + Extract tool names from the request body for the given route. + Uses guardrail translation handlers when available, else standalone extractors + for generate_content and MCP. Returns [] for non-tool-capable routes or when + no tools are present. + """ + call_types = get_call_types_for_route(route) + if not call_types: + return [] + _register_standalone_extractors() + mappings = load_guardrail_translation_mappings() + for call_type in call_types: + if not isinstance(call_type, CallTypes): + continue + if call_type.value not in TOOL_CAPABLE_CALL_TYPES: + continue + if call_type.value in STANDALONE_EXTRACTORS: + return STANDALONE_EXTRACTORS[call_type.value](data) + handler_cls = mappings.get(call_type) + if handler_cls is not None: + names = handler_cls().extract_request_tool_names(data) + if names: + return names + return [] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 92bf035a986..32eab99fb99 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -89,6 +89,25 @@ def _get_metadata_variable_name(request: Request) -> str: return "metadata" +def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str]: + """ + Extract chain id for call chaining from request headers. + + x-litellm-trace-id and x-litellm-session-id are interchangeable; when both + are present, x-litellm-trace-id takes precedence. Header keys are matched + case-insensitively so this works with raw header dicts from any transport. + + Used by MCP (and other paths that have raw_headers but no Request) to set + litellm_trace_id/litellm_session_id for spend logs and logging consistency. + """ + if not headers: + return None + normalized = {k.lower(): v for k, v in headers.items() if isinstance(k, str)} + return normalized.get("x-litellm-trace-id") or normalized.get( + "x-litellm-session-id" + ) + + def safe_add_api_version_from_query_params(data: dict, request: Request): try: if hasattr(request, "query_params"): @@ -177,12 +196,12 @@ def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> Optional[TeamCallbackMetadata]: callback_settings_obj: Optional[TeamCallbackMetadata] = None - key_dynamic_logging_settings: Optional[ - dict - ] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) - team_dynamic_logging_settings: Optional[ - dict - ] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) + key_dynamic_logging_settings: Optional[dict] = ( + KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + ) + team_dynamic_logging_settings: Optional[dict] = ( + KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) + ) ######################################################################################### # Key-based callbacks ######################################################################################### @@ -576,9 +595,13 @@ class LiteLLMProxyRequestSetup: ######################################################################################### # Finally update the requests metadata with the `metadata_from_headers` ######################################################################################### + agent_id_from_header = headers.get("x-litellm-agent-id") - trace_id_from_header = headers.get("x-litellm-trace-id") - session_id_from_header = headers.get("x-litellm-session-id") + # x-litellm-trace-id and x-litellm-session-id are interchangeable for call chaining + chain_id = headers.get("x-litellm-trace-id") or headers.get( + "x-litellm-session-id" + ) + if agent_id_from_header: metadata_from_headers["agent_id"] = agent_id_from_header @@ -586,16 +609,13 @@ class LiteLLMProxyRequestSetup: f"Extracted agent_id from header: {agent_id_from_header}" ) - if trace_id_from_header: - metadata_from_headers["trace_id"] = trace_id_from_header + if chain_id: + metadata_from_headers["trace_id"] = chain_id + metadata_from_headers["session_id"] = chain_id + data["litellm_session_id"] = chain_id + data["litellm_trace_id"] = chain_id verbose_proxy_logger.debug( - f"Extracted trace_id from header: {trace_id_from_header}" - ) - - if session_id_from_header: - metadata_from_headers["session_id"] = session_id_from_header - verbose_proxy_logger.debug( - f"Extracted session_id from header: {session_id_from_header}" + f"Extracted chain_id from header (trace-id/session-id): {chain_id}" ) if isinstance(data[_metadata_variable_name], dict): @@ -702,11 +722,11 @@ class LiteLLMProxyRequestSetup: ## KEY-LEVEL SPEND LOGS / TAGS if "tags" in key_metadata and key_metadata["tags"] is not None: - data[_metadata_variable_name][ - "tags" - ] = LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=key_metadata["tags"], + data[_metadata_variable_name]["tags"] = ( + LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=key_metadata["tags"], + ) ) if "disable_global_guardrails" in key_metadata and isinstance( key_metadata["disable_global_guardrails"], bool @@ -839,14 +859,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 """ from litellm.proxy.proxy_server import llm_router, premium_user - from litellm.types.proxy.litellm_pre_call_utils import ( - RedactedDict, - SecretFields, - ) + from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields - _raw_headers: Dict[str, str] = RedactedDict( - _safe_get_request_headers(request) - ) + _raw_headers: Dict[str, str] = RedactedDict(_safe_get_request_headers(request)) forward_llm_auth = False if general_settings: @@ -986,9 +1001,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data[_metadata_variable_name]["litellm_api_version"] = version if general_settings is not None: - data[_metadata_variable_name][ - "global_max_parallel_requests" - ] = general_settings.get("global_max_parallel_requests", None) + data[_metadata_variable_name]["global_max_parallel_requests"] = ( + general_settings.get("global_max_parallel_requests", None) + ) ### KEY-LEVEL Controls key_metadata = user_api_key_dict.metadata @@ -1076,6 +1091,15 @@ async def add_litellm_data_to_request( # noqa: PLR0915 ] = user_api_key_dict.user_max_budget data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata + data[_metadata_variable_name]["user_api_key_team_metadata"] = ( + user_api_key_dict.team_metadata + ) + data[_metadata_variable_name]["user_api_key_object_permission_id"] = ( + getattr(user_api_key_dict, "object_permission_id", None) + ) + data[_metadata_variable_name]["user_api_key_team_object_permission_id"] = ( + getattr(user_api_key_dict, "team_object_permission_id", None) + ) data[_metadata_variable_name]["headers"] = _headers data[_metadata_variable_name]["endpoint"] = str(request.url) diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 89880c9a4ec..7fdd3475c04 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -4,27 +4,87 @@ TOOL POLICY MANAGEMENT All /tool management endpoints GET /v1/tool/list - List all discovered tools and their policies +GET /v1/tool/policy/options - List available input/output policy options with descriptions GET /v1/tool/{tool_name} - Get a single tool's details -POST /v1/tool/policy - Update the call_policy for a tool +POST /v1/tool/policy - Update the input_policy / output_policy for a tool """ -from typing import Optional +import uuid +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, List, Optional -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Query + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.tool_management import ( LiteLLM_ToolTableRow, - ToolCallPolicy, + ToolDetailResponse, + ToolInputPolicy, ToolListResponse, + ToolOutputPolicy, + ToolPolicyOption, + ToolPolicyOptionsResponse, ToolPolicyUpdateRequest, ToolPolicyUpdateResponse, + ToolUsageLogEntry, + ToolUsageLogsResponse, ) router = APIRouter() +TOOL_POLICY_OPTIONS = ToolPolicyOptionsResponse( + input_policies=[ + ToolPolicyOption( + value="untrusted", + label="Untrusted", + description="Tool accepts any input, including data from untrusted tool outputs. Default for newly discovered tools.", + ), + ToolPolicyOption( + value="trusted", + label="Trusted", + description="Tool requires trusted input. Blocked if the conversation contains output from any tool with output_policy=untrusted.", + ), + ToolPolicyOption( + value="blocked", + label="Blocked", + description="Tool is completely prohibited. Any attempt to call it is rejected.", + ), + ], + output_policies=[ + ToolPolicyOption( + value="untrusted", + label="Untrusted", + description="Tool output may contain unsafe content (prompt injection, risky code). Downstream tools with input_policy=trusted will be blocked.", + ), + ToolPolicyOption( + value="trusted", + label="Trusted", + description="Tool output is verified safe. Will not trigger trust-chain blocks on downstream tools.", + ), + ], +) + + +@router.get( + "/v1/tool/policy/options", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolPolicyOptionsResponse, +) +async def get_tool_policy_options( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Return the available input and output policy options with descriptions. + Static data — no DB call. + """ + return TOOL_POLICY_OPTIONS + @router.get( "/v1/tool/list", @@ -33,14 +93,14 @@ router = APIRouter() response_model=ToolListResponse, ) async def list_tools( - call_policy: Optional[ToolCallPolicy] = None, + input_policy: Optional[ToolInputPolicy] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - List all auto-discovered tools and their call policies. + List all auto-discovered tools and their policies. Parameters: - - call_policy: Optional filter — one of "trusted", "untrusted", "dual_llm", "blocked" + - input_policy: Optional filter — one of "trusted", "untrusted", "blocked" """ from litellm.proxy.db.tool_registry_writer import list_tools as db_list_tools from litellm.proxy.proxy_server import prisma_client @@ -51,13 +111,201 @@ async def list_tools( ) try: - tools = await db_list_tools(prisma_client=prisma_client, call_policy=call_policy) + tools = await db_list_tools( + prisma_client=prisma_client, input_policy=input_policy + ) return ToolListResponse(tools=tools, total=len(tools)) except Exception as e: verbose_proxy_logger.exception("Error listing tools: %s", e) raise HTTPException(status_code=500, detail=str(e)) +@router.get( + "/v1/tool/{tool_name:path}/detail", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolDetailResponse, +) +async def get_tool_detail( + tool_name: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get a single tool with its policy overrides (for UI detail view). + """ + from litellm.proxy.db.tool_registry_writer import get_tool as db_get_tool + from litellm.proxy.db.tool_registry_writer import list_overrides_for_tool + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name) + if tool is None: + raise HTTPException(status_code=404, detail=f"Tool '{tool_name}' not found") + overrides = await list_overrides_for_tool( + prisma_client=prisma_client, tool_name=tool_name + ) + return ToolDetailResponse(tool=tool, overrides=overrides) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error getting tool detail: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> Optional[str]: + """Short snippet from messages or proxy_server_request for tool usage log row.""" + if sl is None: + return None + messages = getattr(sl, "messages", None) + if messages is not None: + s = _snippet_str(messages, max_len) + if s: + return s + psr = getattr(sl, "proxy_server_request", None) + if not psr: + return None + if isinstance(psr, str): + import json + + try: + psr = json.loads(psr) + except Exception: + return _snippet_str(psr, max_len) + if isinstance(psr, dict): + msgs = psr.get("messages") + if msgs is None and isinstance(psr.get("body"), dict): + msgs = psr["body"].get("messages") + s = _snippet_str(msgs, max_len) + if s: + return s + return _snippet_str(psr, max_len) + + +def _snippet_str(text: Any, max_len: int = 200) -> Optional[str]: + if text is None: + return None + if isinstance(text, str): + s = text + elif isinstance(text, list): + parts = [] + for item in text: + if isinstance(item, dict) and "content" in item: + c = item["content"] + parts.append(c if isinstance(c, str) else str(c)) + else: + parts.append(str(item)) + s = " ".join(parts) + else: + s = str(text) + if not s or s == "{}": + return None + return (s[:max_len] + "...") if len(s) > max_len else s + + +@router.get( + "/v1/tool/{tool_name:path}/logs", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolUsageLogsResponse, +) +async def get_tool_usage_logs( + tool_name: str, + page: int = Query(1, ge=1), + page_size: int = Query(50, ge=1, le=100), + start_date: Optional[str] = Query(None, description="YYYY-MM-DD"), + end_date: Optional[str] = Query(None, description="YYYY-MM-DD"), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Return paginated spend logs for requests that used this tool (from SpendLogToolIndex). + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + where: dict = {"tool_name": tool_name} + if start_date or end_date: + start_time_filter: Optional[datetime] = None + end_time_filter: Optional[datetime] = None + if start_date: + try: + start_time_filter = datetime.strptime( + start_date + "T00:00:00", "%Y-%m-%dT%H:%M:%S" + ).replace(tzinfo=timezone.utc) + except ValueError: + pass + if end_date: + try: + end_time_filter = datetime.strptime( + end_date + "T23:59:59", "%Y-%m-%dT%H:%M:%S" + ).replace(tzinfo=timezone.utc) + except ValueError: + pass + if start_time_filter is not None or end_time_filter is not None: + where["start_time"] = {} + if start_time_filter is not None: + where["start_time"]["gte"] = start_time_filter + if end_time_filter is not None: + where["start_time"]["lte"] = end_time_filter + + total = await prisma_client.db.litellm_spendlogtoolindex.count(where=where) + index_rows = await prisma_client.db.litellm_spendlogtoolindex.find_many( + where=where, + order={"start_time": "desc"}, + skip=(page - 1) * page_size, + take=page_size, + ) + request_ids = [r.request_id for r in index_rows] + if not request_ids: + return ToolUsageLogsResponse( + logs=[], total=total, page=page, page_size=page_size + ) + + spend_logs = await prisma_client.db.litellm_spendlogs.find_many( + where={"request_id": {"in": request_ids}} + ) + log_by_id = {s.request_id: s for s in spend_logs} + + logs_out: List[ToolUsageLogEntry] = [] + for r in index_rows: + sl = log_by_id.get(r.request_id) + if not sl: + continue + ts = ( + sl.startTime.isoformat() + if hasattr(sl.startTime, "isoformat") + else str(sl.startTime) + ) + logs_out.append( + ToolUsageLogEntry( + id=sl.request_id, + timestamp=ts, + model=getattr(sl, "model", None) or None, + spend=getattr(sl, "spend", None), + total_tokens=getattr(sl, "total_tokens", None), + input_snippet=_input_snippet_for_tool_log(sl), + ) + ) + + return ToolUsageLogsResponse( + logs=logs_out, total=total, page=page, page_size=page_size + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error getting tool usage logs: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + @router.get( "/v1/tool/{tool_name:path}", tags=["tool management"], @@ -70,9 +318,6 @@ async def get_tool( ): """ Get details for a single tool. - - Parameters: - - tool_name: The tool name (supports namespaced names with slashes) """ from litellm.proxy.db.tool_registry_writer import get_tool as db_get_tool from litellm.proxy.proxy_server import prisma_client @@ -85,9 +330,7 @@ async def get_tool( try: tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name) if tool is None: - raise HTTPException( - status_code=404, detail=f"Tool '{tool_name}' not found" - ) + raise HTTPException(status_code=404, detail=f"Tool '{tool_name}' not found") return tool except HTTPException: raise @@ -96,6 +339,80 @@ async def get_tool( raise HTTPException(status_code=500, detail=str(e)) +async def _resolve_key_hash_to_object_permission_id( + prisma_client: "PrismaClient", + key_hash: str, +) -> Optional[str]: + """Resolve key (hash or raw) to object_permission_id; create permission if key has none.""" + from litellm.proxy.proxy_server import hash_token + + hashed = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) + if not hashed: + return None + row = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed} + ) + if row is None: + return None + op_id = getattr(row, "object_permission_id", None) + if op_id: + return op_id + new_id = str(uuid.uuid4()) + await prisma_client.db.litellm_objectpermissiontable.create( + data={"object_permission_id": new_id, "blocked_tools": []} + ) + updated_count = await prisma_client.db.litellm_verificationtoken.update_many( + where={"token": hashed, "object_permission_id": None}, + data={"object_permission_id": new_id}, + ) + if updated_count == 0: + await prisma_client.db.litellm_objectpermissiontable.delete( + where={"object_permission_id": new_id} + ) + row = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed} + ) + return getattr(row, "object_permission_id", None) if row else None + return new_id + + +async def _resolve_team_id_to_object_permission_id( + prisma_client: "PrismaClient", + team_id: str, +) -> Optional[str]: + """Resolve team_id to object_permission_id; create permission if team has none.""" + if not team_id or not team_id.strip(): + return None + team_id_clean = team_id.strip() + row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id_clean}, + select={"object_permission_id": True}, + ) + if row is None: + return None + op_id = getattr(row, "object_permission_id", None) + if op_id: + return op_id + new_id = str(uuid.uuid4()) + await prisma_client.db.litellm_objectpermissiontable.create( + data={"object_permission_id": new_id, "blocked_tools": []} + ) + updated_count = await prisma_client.db.litellm_teamtable.update_many( + where={"team_id": team_id_clean, "object_permission_id": None}, + data={"object_permission_id": new_id}, + ) + if updated_count == 0: + await prisma_client.db.litellm_objectpermissiontable.delete( + where={"object_permission_id": new_id} + ) + row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id_clean}, + select={"object_permission_id": True}, + ) + return getattr(row, "object_permission_id", None) if row else None + return new_id + + @router.post( "/v1/tool/policy", tags=["tool management"], @@ -107,15 +424,20 @@ async def update_tool_policy( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Set the call policy for a tool. + Set the input_policy and/or output_policy for a tool (global), or block for a specific team/key (override). Parameters: - tool_name: str - The tool to update - - call_policy: "trusted" | "untrusted" | "dual_llm" | "blocked" - - Setting a tool to "blocked" will cause the ToolPolicyGuardrail to remove - that tool_call from LLM responses before returning them to the client. + - input_policy: optional - "trusted" | "untrusted" | "blocked" + - output_policy: optional - "trusted" | "untrusted" + - team_id: optional - if set, create/update override for this team only + - key_hash: optional - if set, create/update override for this key only """ + from litellm.proxy.db.tool_registry_writer import ( + add_tool_to_object_permission_blocked, + get_tool_policy_registry, + remove_tool_from_object_permission_blocked, + ) from litellm.proxy.db.tool_registry_writer import ( update_tool_policy as db_update_tool_policy, ) @@ -127,19 +449,80 @@ async def update_tool_policy( ) try: + if data.team_id is not None or data.key_hash is not None: + if data.team_id is not None and data.key_hash is not None: + raise HTTPException( + status_code=400, + detail="Provide either team_id or key_hash, not both", + ) + if data.key_hash is not None: + op_id = await _resolve_key_hash_to_object_permission_id( + prisma_client, data.key_hash + ) + else: + op_id = await _resolve_team_id_to_object_permission_id( + prisma_client, data.team_id or "" + ) + if op_id is None: + raise HTTPException( + status_code=404, + detail="Key or team not found for the given identifier", + ) + is_blocking = data.input_policy == "blocked" + if is_blocking: + ok = await add_tool_to_object_permission_blocked( + prisma_client=prisma_client, + object_permission_id=op_id, + tool_name=data.tool_name, + ) + else: + ok = await remove_tool_from_object_permission_blocked( + prisma_client=prisma_client, + object_permission_id=op_id, + tool_name=data.tool_name, + ) + if not ok: + raise HTTPException( + status_code=500, + detail=f"Failed to update policy override for tool '{data.tool_name}'", + ) + registry = get_tool_policy_registry() + if registry.is_initialized(): + await registry.sync_tool_policy_from_db(prisma_client) + return ToolPolicyUpdateResponse( + tool_name=data.tool_name, + input_policy=data.input_policy, + output_policy=data.output_policy, + updated=True, + team_id=data.team_id, + key_hash=data.key_hash, + ) + + if data.input_policy is None and data.output_policy is None: + raise HTTPException( + status_code=400, + detail="At least one of input_policy or output_policy must be provided", + ) + updated = await db_update_tool_policy( prisma_client=prisma_client, tool_name=data.tool_name, - call_policy=data.call_policy, updated_by=user_api_key_dict.user_id, + input_policy=data.input_policy, + output_policy=data.output_policy, ) if updated is None: raise HTTPException( - status_code=500, detail=f"Failed to update policy for tool '{data.tool_name}'" + status_code=500, + detail=f"Failed to update policy for tool '{data.tool_name}'", ) + registry = get_tool_policy_registry() + if registry.is_initialized(): + await registry.sync_tool_policy_from_db(prisma_client) return ToolPolicyUpdateResponse( tool_name=updated.tool_name, - call_policy=updated.call_policy, + input_policy=updated.input_policy, + output_policy=updated.output_policy, updated=True, ) except HTTPException: @@ -147,3 +530,77 @@ async def update_tool_policy( except Exception as e: verbose_proxy_logger.exception("Error updating tool policy: %s", e) raise HTTPException(status_code=500, detail=str(e)) + + +@router.delete( + "/v1/tool/{tool_name:path}/overrides", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_tool_policy_override( + tool_name: str, + team_id: Optional[str] = Query( + None, description="Team ID of the override to remove" + ), + key_hash: Optional[str] = Query( + None, description="Key hash of the override to remove" + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Remove a policy override for a tool. Specify the override by team_id or key_hash + (exactly one required). + """ + from litellm.proxy.db.tool_registry_writer import ( + get_tool_policy_registry, + remove_tool_from_object_permission_blocked, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + if team_id is None and key_hash is None: + raise HTTPException( + status_code=400, + detail="At least one of team_id or key_hash is required to identify the override", + ) + if team_id is not None and key_hash is not None: + raise HTTPException( + status_code=400, + detail="Provide either team_id or key_hash, not both", + ) + try: + if key_hash is not None: + op_id = await _resolve_key_hash_to_object_permission_id( + prisma_client, key_hash + ) + else: + op_id = await _resolve_team_id_to_object_permission_id( + prisma_client, team_id or "" + ) + if op_id is None: + raise HTTPException( + status_code=404, + detail="Key or team not found for the given identifier", + ) + deleted = await remove_tool_from_object_permission_blocked( + prisma_client=prisma_client, + object_permission_id=op_id, + tool_name=tool_name, + ) + if not deleted: + raise HTTPException( + status_code=404, + detail=f"No override found for tool '{tool_name}' with the given scope", + ) + registry = get_tool_policy_registry() + if registry.is_initialized(): + await registry.sync_tool_policy_from_db(prisma_client) + return {"deleted": True, "tool_name": tool_name} + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error deleting tool policy override: %s", e) + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b3d707b1aa2..6a2b0accb0e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4411,6 +4411,9 @@ class ProxyConfig: if self._should_load_db_object(object_type="search_tools"): await self._init_search_tools_in_db(prisma_client=prisma_client) + if self._should_load_db_object(object_type="tools"): + await self._init_tool_policy_in_db(prisma_client=prisma_client) + if self._should_load_db_object(object_type="model_cost_map"): await self._check_and_reload_model_cost_map(prisma_client=prisma_client) @@ -4847,6 +4850,24 @@ class ProxyConfig: ) ) + async def _init_tool_policy_in_db(self, prisma_client: PrismaClient): + """ + Initialize tool policy from database into the in-memory registry. + Synced periodically by add_deployment -> _init_non_llm_objects_in_db. + """ + from litellm.proxy.db.tool_registry_writer import get_tool_policy_registry + + try: + registry = get_tool_policy_registry() + await registry.sync_tool_policy_from_db(prisma_client=prisma_client) + verbose_proxy_logger.debug("Successfully synced tool policy from DB") + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_tool_policy_in_db - {}".format( + str(e) + ) + ) + async def _init_vector_stores_in_db(self, prisma_client: PrismaClient): from litellm.vector_stores.vector_store_registry import VectorStoreRegistry @@ -10577,6 +10598,12 @@ async def async_queue_request( data["metadata"]["user_api_key_team_id"] = getattr( user_api_key_dict, "team_id", None ) + data["metadata"]["user_api_key_object_permission_id"] = getattr( + user_api_key_dict, "object_permission_id", None + ) + data["metadata"]["user_api_key_team_object_permission_id"] = getattr( + user_api_key_dict, "team_object_permission_id", None + ) data["metadata"]["endpoint"] = str(request.url) global user_temperature, user_request_timeout, user_max_tokens, user_api_base @@ -11093,9 +11120,7 @@ async def get_favicon(): if favicon_url.startswith(("http://", "https://")): try: - from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, - ) + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider async_client = get_async_httpx_client( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index e0b28a4e012..25ee2750548 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -260,6 +260,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -928,6 +929,16 @@ model LiteLLM_SpendLogGuardrailIndex { @@index([policy_id, start_time]) } +// Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production +model LiteLLM_SpendLogToolIndex { + request_id String + tool_name String // matches LiteLLM_ToolTable.tool_name; join for input_policy/output_policy etc. + start_time DateTime + + @@id([request_id, tool_name]) + @@index([tool_name, start_time]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) @@ -1065,23 +1076,27 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } -// Global tool registry - auto-discovered from LLM responses; admins set call_policy here +// Global tool registry - auto-discovered from LLM responses; admins set input/output policies here model LiteLLM_ToolTable { - tool_id String @id @default(uuid()) - tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" - origin String? // MCP server name or "user_defined" - call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" - call_count Int @default(0) // cumulative number of times this tool was seen - assignments Json? @default("{}") - key_hash String? // hash of the virtual key that first called this tool - team_id String? // team that first called this tool - key_alias String? // human-readable alias of the virtual key - created_at DateTime @default(now()) - created_by String? - updated_at DateTime @default(now()) @updatedAt - updated_by String? + tool_id String @id @default(uuid()) + tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" + origin String? // MCP server name or "user_defined" + input_policy String @default("untrusted") // "trusted" | "untrusted" | "blocked" + output_policy String @default("untrusted") // "trusted" | "untrusted" + call_count Int @default(0) // cumulative number of times this tool was seen + assignments Json? @default("{}") + key_hash String? // hash of the virtual key that first called this tool + team_id String? // team that first called this tool + key_alias String? // human-readable alias of the virtual key + user_agent String? // user-agent of the first request that discovered this tool + last_used_at DateTime? // timestamp of the most recent call + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? - @@index([call_policy]) + @@index([input_policy]) + @@index([output_policy]) @@index([team_id]) } diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 31615a768d7..131841f7b59 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,26 +11,21 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import ( - MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, -) +from litellm.constants import \ + MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.core_helpers import ( - get_litellm_metadata_from_kwargs, - reconstruct_model_name, -) + get_litellm_metadata_from_kwargs, reconstruct_model_name) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.utils import PrismaClient, hash_token -from litellm.types.utils import ( - CostBreakdown, - StandardLoggingGuardrailInformation, - StandardLoggingMCPToolCall, - StandardLoggingModelInformation, - StandardLoggingPayload, - StandardLoggingVectorStoreRequest, - VectorStoreSearchResponse, -) +from litellm.types.utils import (CostBreakdown, + StandardLoggingGuardrailInformation, + StandardLoggingMCPToolCall, + StandardLoggingModelInformation, + StandardLoggingPayload, + StandardLoggingVectorStoreRequest, + VectorStoreSearchResponse) from litellm.utils import get_end_user_id_for_cost_tracking @@ -116,16 +111,15 @@ def _get_spend_logs_metadata( # Filter the metadata dictionary to include only the specified keys clean_metadata = SpendLogsMetadata( **{ # type: ignore - key: metadata.get(key) - for key in SpendLogsMetadata.__annotations__.keys() + key: metadata.get(key) for key in SpendLogsMetadata.__annotations__.keys() } ) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata - clean_metadata[ - "vector_store_request_metadata" - ] = _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) + clean_metadata["vector_store_request_metadata"] = ( + _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) + ) clean_metadata["guardrail_information"] = guardrail_information clean_metadata["usage_object"] = usage_object clean_metadata["model_map_information"] = model_map_information @@ -372,9 +366,11 @@ def get_logging_payload( # noqa: PLR0915 guardrail_information=( standard_logging_payload.get("guardrail_information", None) if standard_logging_payload is not None - else metadata.get("standard_logging_guardrail_information", None) - if metadata is not None - else None + else ( + metadata.get("standard_logging_guardrail_information", None) + if metadata is not None + else None + ) ), cold_storage_object_key=( standard_logging_payload["metadata"].get("cold_storage_object_key", None) @@ -501,6 +497,7 @@ def _get_session_id_for_spend_log( """ from litellm._uuid import uuid + if ( standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None @@ -515,9 +512,7 @@ def _get_session_id_for_spend_log( return str(uuid.uuid4()) -def _get_request_duration_ms( - start_time: datetime, end_time: datetime -) -> Optional[int]: +def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> Optional[int]: """Compute request duration in milliseconds from start and end times.""" try: return int((end_time - start_time).total_seconds() * 1000) @@ -709,20 +704,20 @@ def _convert_to_json_serializable_dict( if max_depth <= 0: # Return a placeholder if max depth is exceeded return "" - + if visited is None: visited = set() - + # Get the object's memory address to track visited objects obj_id = id(obj) if obj_id in visited: # Circular reference detected, return placeholder return "" - + # Only track mutable objects (dict, list, objects with __dict__) if isinstance(obj, (dict, list)) or hasattr(obj, "__dict__"): visited.add(obj_id) - + try: if isinstance(obj, BaseModel): # Use Pydantic's model_dump() instead of pickle @@ -741,7 +736,9 @@ def _convert_to_json_serializable_dict( ] elif hasattr(obj, "__dict__"): # Handle objects with __dict__ attribute - return _convert_to_json_serializable_dict(obj.__dict__, visited, max_depth - 1) + return _convert_to_json_serializable_dict( + obj.__dict__, visited, max_depth - 1 + ) else: # Primitives (str, int, float, bool, None) pass through return obj @@ -777,9 +774,7 @@ def _get_proxy_server_request_for_spend_logs_payload( # Apply message redaction if turn_off_message_logging is enabled if kwargs is not None: from litellm.litellm_core_utils.redact_messages import ( - perform_redaction, - should_redact_message_logging, - ) + perform_redaction, should_redact_message_logging) # Build model_call_details dict to check redaction settings model_call_details = { @@ -788,12 +783,12 @@ def _get_proxy_server_request_for_spend_logs_payload( "standard_callback_dynamic_params" ), } - + # If redaction is enabled, convert to serializable dict before redacting if should_redact_message_logging(model_call_details=model_call_details): _request_body = _convert_to_json_serializable_dict(_request_body) perform_redaction(model_call_details=_request_body, result=None) - + _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) _request_body_json_str = json.dumps(_request_body, default=str) return _request_body_json_str @@ -845,10 +840,8 @@ def _get_response_for_spend_logs_payload( # Apply message redaction if turn_off_message_logging is enabled if kwargs is not None: from litellm.litellm_core_utils.redact_messages import ( - perform_redaction, - should_redact_message_logging, - ) - + perform_redaction, should_redact_message_logging) + litellm_params = kwargs.get("litellm_params", {}) model_call_details = { "litellm_params": litellm_params, @@ -856,11 +849,13 @@ def _get_response_for_spend_logs_payload( "standard_callback_dynamic_params" ), } - + # If redaction is enabled, convert to serializable dict before redacting if should_redact_message_logging(model_call_details=model_call_details): response_obj = _convert_to_json_serializable_dict(response_obj) - response_obj = perform_redaction(model_call_details={}, result=response_obj) + response_obj = perform_redaction( + model_call_details={}, result=response_obj + ) sanitized_wrapper = _sanitize_request_body_for_spend_logs_payload( {"response": response_obj} @@ -882,7 +877,7 @@ def _should_store_prompts_and_responses_in_spend_logs() -> bool: # Check general_settings (from DB or proxy_config.yaml) store_prompts_value = general_settings.get("store_prompts_in_spend_logs") - + # Normalize case: handle True/true/TRUE, False/false/FALSE, None/null if store_prompts_value is True: return True @@ -890,7 +885,7 @@ def _should_store_prompts_and_responses_in_spend_logs() -> bool: # Case-insensitive string comparison if store_prompts_value.lower() == "true": return True - + # Also check environment variable return get_secret_bool("STORE_PROMPTS_IN_SPEND_LOGS") is True diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index afcdd9d0c50..e6da95bb78f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3583,8 +3583,9 @@ class PrismaClient: def _get_engine_pid(self) -> int: try: engine = self.db._original_prisma._engine # type: ignore[attr-defined] - if engine is not None and engine.process is not None: - return engine.process.pid + process = getattr(engine, "process", None) if engine is not None else None + if process is not None: + return process.pid except (AttributeError, TypeError): pass return 0 @@ -4688,6 +4689,19 @@ async def update_spend_logs_job( guardrail_tracking_err, ) + # Tool usage tracking (same batch): SpendLogToolIndex for "last N requests for tool X" + try: + from litellm.proxy.db.spend_log_tool_index import process_spend_logs_tool_usage + await process_spend_logs_tool_usage( + prisma_client=prisma_client, + logs_to_process=logs_to_process, + ) + except Exception as tool_tracking_err: + verbose_proxy_logger.warning( + "Spend tracking - tool usage tracking failed (non-fatal): %s", + tool_tracking_err, + ) + async def _monitor_spend_logs_queue( prisma_client: PrismaClient, diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index 8704ff27759..1c5e1df9e9a 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -5,21 +5,27 @@ Pydantic models for Tool Policy management endpoints. from datetime import datetime from typing import Dict, List, Literal, Optional -from pydantic import BaseModel +from pydantic import BaseModel, Field ToolCallPolicy = Literal["trusted", "untrusted", "dual_llm", "blocked"] +ToolInputPolicy = Literal["trusted", "untrusted", "blocked"] +ToolOutputPolicy = Literal["trusted", "untrusted"] + class LiteLLM_ToolTableRow(BaseModel): tool_id: str tool_name: str origin: Optional[str] = None - call_policy: ToolCallPolicy = "untrusted" + input_policy: ToolInputPolicy = "untrusted" + output_policy: ToolOutputPolicy = "untrusted" call_count: int = 0 assignments: Optional[Dict] = None key_hash: Optional[str] = None team_id: Optional[str] = None key_alias: Optional[str] = None + user_agent: Optional[str] = None + last_used_at: Optional[datetime] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None created_by: Optional[str] = None @@ -33,10 +39,62 @@ class ToolListResponse(BaseModel): class ToolPolicyUpdateRequest(BaseModel): tool_name: str - call_policy: ToolCallPolicy + input_policy: Optional[ToolInputPolicy] = None + output_policy: Optional[ToolOutputPolicy] = None + team_id: Optional[str] = None + key_hash: Optional[str] = None + key_alias: Optional[str] = None class ToolPolicyUpdateResponse(BaseModel): tool_name: str - call_policy: ToolCallPolicy + input_policy: Optional[ToolInputPolicy] = None + output_policy: Optional[ToolOutputPolicy] = None updated: bool + team_id: Optional[str] = None + key_hash: Optional[str] = None + + +class ToolPolicyOverrideRow(BaseModel): + override_id: str + tool_name: str + team_id: Optional[str] = None + key_hash: Optional[str] = None + input_policy: ToolInputPolicy = "blocked" + key_alias: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + + +class ToolPolicyOption(BaseModel): + value: str + label: str + description: str + + +class ToolPolicyOptionsResponse(BaseModel): + input_policies: List[ToolPolicyOption] + output_policies: List[ToolPolicyOption] + + +class ToolDetailResponse(BaseModel): + tool: LiteLLM_ToolTableRow + overrides: List[ToolPolicyOverrideRow] = Field(default_factory=list) + + +class ToolUsageLogEntry(BaseModel): + """One spend log row for a tool call (for UI "recent logs" table).""" + + id: str # request_id + timestamp: str + model: Optional[str] = None + spend: Optional[float] = None + total_tokens: Optional[int] = None + input_snippet: Optional[str] = None + + +class ToolUsageLogsResponse(BaseModel): + logs: List[ToolUsageLogEntry] + total: int + page: int + page_size: int diff --git a/litellm/utils.py b/litellm/utils.py index d192609eead..375bea724b8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1454,10 +1454,12 @@ def client(original_function): # noqa: PLR0915 logging_obj, kwargs = function_setup( original_function.__name__, rules_obj, start_time, *args, **kwargs ) - + # Type assertion: logging_obj is guaranteed to be non-None after function_setup - assert logging_obj is not None, "logging_obj should not be None after function_setup" - + assert ( + logging_obj is not None + ), "logging_obj should not be None after function_setup" + ## LOAD CREDENTIALS load_credentials_from_list(kwargs) kwargs["litellm_logging_obj"] = logging_obj @@ -1753,7 +1755,9 @@ def client(original_function): # noqa: PLR0915 print_args_passed_to_litellm(original_function, args, kwargs) start_time = datetime.datetime.now() result = None - _update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") + _update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get( "litellm_logging_obj", None ) @@ -1776,9 +1780,11 @@ def client(original_function): # noqa: PLR0915 logging_obj, kwargs = function_setup( original_function.__name__, rules_obj, start_time, *args, **kwargs ) - + # Type assertion: logging_obj is guaranteed to be non-None after function_setup - assert logging_obj is not None, "logging_obj should not be None after function_setup" + assert ( + logging_obj is not None + ), "logging_obj should not be None after function_setup" modified_kwargs = await async_pre_call_deployment_hook(kwargs, call_type) if modified_kwargs is not None: @@ -1861,6 +1867,7 @@ def client(original_function): # noqa: PLR0915 # MODEL CALL result = await original_function(*args, **kwargs) end_time = datetime.datetime.now() + if _is_streaming_request( kwargs=kwargs, call_type=call_type, @@ -2082,12 +2089,14 @@ def _is_async_request( return False -_STREAMING_CALL_TYPES = frozenset({ - CallTypes.generate_content_stream, - CallTypes.agenerate_content_stream, - CallTypes.generate_content_stream.value, - CallTypes.agenerate_content_stream.value, -}) +_STREAMING_CALL_TYPES = frozenset( + { + CallTypes.generate_content_stream, + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream.value, + CallTypes.agenerate_content_stream.value, + } +) def _is_streaming_request( @@ -2181,7 +2190,7 @@ def encode(model="", text="", custom_tokenizer: Optional[dict] = None): # Normalize: HuggingFace Tokenizer.encode() returns an Encoding object; # extract .ids so the return type is always List[int]. if hasattr(enc, "ids"): - return enc.ids + return enc.ids # type: ignore return enc @@ -5836,7 +5845,7 @@ def get_model_info( _model_info[key] = value # type: ignore # if verbose_logger.isEnabledFor(logging.DEBUG): - # verbose_logger.debug(f"model_info: {_model_info}") + # verbose_logger.debug(f"model_info: {_model_info}") returned_model_info = ModelInfo( **_model_info, supported_openai_params=supported_openai_params @@ -6179,8 +6188,10 @@ def validate_environment( # noqa: PLR0915 "AWS_ROLE_ARN" in os.environ or "AWS_PROFILE" in os.environ or "AWS_WEB_IDENTITY_TOKEN_FILE" in os.environ - or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" in os.environ # ECS task role - or "AWS_CONTAINER_CREDENTIALS_FULL_URI" in os.environ # ECS/Fargate full URI credential delivery + or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" + in os.environ # ECS task role + or "AWS_CONTAINER_CREDENTIALS_FULL_URI" + in os.environ # ECS/Fargate full URI credential delivery ): keys_in_environment = True else: @@ -7386,7 +7397,9 @@ class ModelResponseIterator: if convert_to_delta is True: _stream_response = ModelResponseStream() _stream_response.choices[0].delta.content = model_response.choices[0].message.content # type: ignore - self.model_response: Union[ModelResponse, ModelResponseStream] = _stream_response + self.model_response: Union[ModelResponse, ModelResponseStream] = ( + _stream_response + ) else: self.model_response = model_response self.is_done = False @@ -7457,13 +7470,22 @@ def is_cached_message(message: AllMessageValues) -> bool: Used for anthropic/gemini context caching. Follows the anthropic format {"cache_control": {"type": "ephemeral"}} - + Can be disabled globally by setting litellm.disable_anthropic_gemini_context_caching_transform = True """ # Check if context caching is disabled globally if litellm.disable_anthropic_gemini_context_caching_transform is True: return False - + + # Check message-level cache_control (set by cache_control_injection_points hook for string content) + message_level_cache_control = message.get("cache_control") + if ( + message_level_cache_control is not None + and isinstance(message_level_cache_control, dict) + and message_level_cache_control.get("type") == "ephemeral" + ): + return True + if "content" not in message: return False @@ -7980,6 +8002,7 @@ class ProviderConfigManager: def _get_azure_ai_config(model: str) -> BaseConfig: """Get Azure AI config based on model type.""" from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + return AzureFoundryModelInfo.get_azure_ai_config_for_model(model) @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 412f99791a0..b92e2727979 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10960,7 +10960,8 @@ "output_cost_per_token": 9e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": { "max_tokens": 131072, @@ -10970,7 +10971,8 @@ "output_cost_per_token": 1e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": { "max_tokens": 131072, @@ -10990,7 +10992,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen2.5-72B-Instruct": { "max_tokens": 32768, @@ -11000,7 +11003,8 @@ "output_cost_per_token": 3.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen2.5-7B-Instruct": { "max_tokens": 32768, @@ -11021,7 +11025,8 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-14B": { "max_tokens": 40960, @@ -11031,7 +11036,8 @@ "output_cost_per_token": 2.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B": { "max_tokens": 40960, @@ -11041,7 +11047,8 @@ "output_cost_per_token": 5.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": { "max_tokens": 262144, @@ -11051,7 +11058,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": { "max_tokens": 262144, @@ -11061,7 +11069,8 @@ "output_cost_per_token": 2.9e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-30B-A3B": { "max_tokens": 40960, @@ -11071,7 +11080,8 @@ "output_cost_per_token": 2.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-32B": { "max_tokens": 40960, @@ -11081,7 +11091,8 @@ "output_cost_per_token": 2.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": { "max_tokens": 262144, @@ -11091,7 +11102,8 @@ "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": { "max_tokens": 262144, @@ -11101,7 +11113,8 @@ "output_cost_per_token": 1.2e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Next-80B-A3B-Instruct": { "max_tokens": 262144, @@ -11111,7 +11124,8 @@ "output_cost_per_token": 1.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Qwen/Qwen3-Next-80B-A3B-Thinking": { "max_tokens": 262144, @@ -11121,7 +11135,8 @@ "output_cost_per_token": 1.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": { "max_tokens": 8192, @@ -11172,7 +11187,8 @@ "cache_read_input_token_cost": 3.3e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/anthropic/claude-4-opus": { "max_tokens": 200000, @@ -11182,7 +11198,8 @@ "output_cost_per_token": 8.25e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/anthropic/claude-4-sonnet": { "max_tokens": 200000, @@ -11192,7 +11209,8 @@ "output_cost_per_token": 1.65e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1": { "max_tokens": 163840, @@ -11202,7 +11220,8 @@ "output_cost_per_token": 2.4e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528": { "max_tokens": 163840, @@ -11213,7 +11232,8 @@ "cache_read_input_token_cost": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": { "max_tokens": 32768, @@ -11223,7 +11243,8 @@ "output_cost_per_token": 3e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { "max_tokens": 131072, @@ -11243,7 +11264,8 @@ "output_cost_per_token": 2.7e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": { "max_tokens": 40960, @@ -11253,7 +11275,8 @@ "output_cost_per_token": 3e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3": { "max_tokens": 163840, @@ -11263,7 +11286,8 @@ "output_cost_per_token": 8.9e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3-0324": { "max_tokens": 163840, @@ -11273,7 +11297,8 @@ "output_cost_per_token": 8.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1": { "max_tokens": 163840, @@ -11285,7 +11310,8 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_reasoning": true + "supports_reasoning": true, + "supports_function_calling": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1-Terminus": { "max_tokens": 163840, @@ -11296,7 +11322,8 @@ "cache_read_input_token_cost": 2.16e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.0-flash-001": { "deprecation_date": "2026-06-01", @@ -11307,7 +11334,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.5-flash": { "max_tokens": 1000000, @@ -11317,7 +11345,8 @@ "output_cost_per_token": 2.5e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -11327,7 +11356,8 @@ "output_cost_per_token": 1e-05, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -11337,7 +11367,8 @@ "output_cost_per_token": 1e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-27b-it": { "max_tokens": 131072, @@ -11347,7 +11378,8 @@ "output_cost_per_token": 1.6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/google/gemma-3-4b-it": { "max_tokens": 131072, @@ -11357,7 +11389,8 @@ "output_cost_per_token": 8e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": { "max_tokens": 131072, @@ -11377,7 +11410,8 @@ "output_cost_per_token": 2e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct": { "max_tokens": 131072, @@ -11387,7 +11421,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": { "max_tokens": 131072, @@ -11397,6 +11432,7 @@ "output_cost_per_token": 3.9e-07, "litellm_provider": "deepinfra", "mode": "chat", + "supports_function_calling": true, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { @@ -11407,7 +11443,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": { "max_tokens": 327680, @@ -11417,7 +11454,8 @@ "output_cost_per_token": 3e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Llama-Guard-3-8B": { "max_tokens": 131072, @@ -11447,7 +11485,8 @@ "output_cost_per_token": 6e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": { "max_tokens": 131072, @@ -11457,7 +11496,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { "max_tokens": 131072, @@ -11467,7 +11507,8 @@ "output_cost_per_token": 2.8e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": { "max_tokens": 131072, @@ -11477,7 +11518,8 @@ "output_cost_per_token": 5e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { "max_tokens": 131072, @@ -11487,7 +11529,8 @@ "output_cost_per_token": 3e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/microsoft/WizardLM-2-8x22B": { "max_tokens": 65536, @@ -11507,7 +11550,8 @@ "output_cost_per_token": 1.4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": { "max_tokens": 131072, @@ -11517,7 +11561,8 @@ "output_cost_per_token": 4e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": { "max_tokens": 32768, @@ -11527,7 +11572,8 @@ "output_cost_per_token": 8e-08, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": { "max_tokens": 128000, @@ -11537,7 +11583,8 @@ "output_cost_per_token": 2e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { "max_tokens": 32768, @@ -11547,7 +11594,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/moonshotai/Kimi-K2-Instruct": { "max_tokens": 131072, @@ -11557,7 +11605,8 @@ "output_cost_per_token": 2e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/moonshotai/Kimi-K2-Instruct-0905": { "max_tokens": 262144, @@ -11568,7 +11617,8 @@ "cache_read_input_token_cost": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": { "max_tokens": 131072, @@ -11578,7 +11628,8 @@ "output_cost_per_token": 6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/Llama-3.3-Nemotron-Super-49B-v1.5": { "max_tokens": 131072, @@ -11588,7 +11639,8 @@ "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { "max_tokens": 131072, @@ -11598,7 +11650,8 @@ "output_cost_per_token": 1.6e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/openai/gpt-oss-120b": { "max_tokens": 131072, @@ -11608,7 +11661,8 @@ "output_cost_per_token": 4.5e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/openai/gpt-oss-20b": { "max_tokens": 131072, @@ -11618,7 +11672,8 @@ "output_cost_per_token": 1.5e-07, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepinfra/zai-org/GLM-4.5": { "max_tokens": 131072, @@ -11628,7 +11683,8 @@ "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", "mode": "chat", - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_function_calling": true }, "deepseek/deepseek-chat": { "cache_creation_input_token_cost": 0.0, @@ -29958,6 +30014,18 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/Qwen/Qwen3.5-397B-A17B": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", diff --git a/schema.prisma b/schema.prisma index e0b28a4e012..25ee2750548 100644 --- a/schema.prisma +++ b/schema.prisma @@ -260,6 +260,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -928,6 +929,16 @@ model LiteLLM_SpendLogGuardrailIndex { @@index([policy_id, start_time]) } +// Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production +model LiteLLM_SpendLogToolIndex { + request_id String + tool_name String // matches LiteLLM_ToolTable.tool_name; join for input_policy/output_policy etc. + start_time DateTime + + @@id([request_id, tool_name]) + @@index([tool_name, start_time]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) @@ -1065,23 +1076,27 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } -// Global tool registry - auto-discovered from LLM responses; admins set call_policy here +// Global tool registry - auto-discovered from LLM responses; admins set input/output policies here model LiteLLM_ToolTable { - tool_id String @id @default(uuid()) - tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" - origin String? // MCP server name or "user_defined" - call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" - call_count Int @default(0) // cumulative number of times this tool was seen - assignments Json? @default("{}") - key_hash String? // hash of the virtual key that first called this tool - team_id String? // team that first called this tool - key_alias String? // human-readable alias of the virtual key - created_at DateTime @default(now()) - created_by String? - updated_at DateTime @default(now()) @updatedAt - updated_by String? + tool_id String @id @default(uuid()) + tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" + origin String? // MCP server name or "user_defined" + input_policy String @default("untrusted") // "trusted" | "untrusted" | "blocked" + output_policy String @default("untrusted") // "trusted" | "untrusted" + call_count Int @default(0) // cumulative number of times this tool was seen + assignments Json? @default("{}") + key_hash String? // hash of the virtual key that first called this tool + team_id String? // team that first called this tool + key_alias String? // human-readable alias of the virtual key + user_agent String? // user-agent of the first request that discovered this tool + last_used_at DateTime? // timestamp of the most recent call + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? - @@index([call_policy]) + @@index([input_policy]) + @@index([output_policy]) @@index([team_id]) } diff --git a/scripts/test_tool_allowlist_script.py b/scripts/test_tool_allowlist_script.py new file mode 100644 index 00000000000..75a50d09b84 --- /dev/null +++ b/scripts/test_tool_allowlist_script.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +""" +Standalone script to test tool allowlist enforcement and tool name extraction. + +Run from repo root: + poetry run python scripts/test_tool_allowlist_script.py + +Or run the unit tests: + poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v +""" + +import asyncio +import sys +from pathlib import Path + +# Ensure repo root is on path +repo_root = Path(__file__).resolve().parent.parent +if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + +def test_extraction(): + """Test extract_request_tool_names for each API shape.""" + from litellm.proxy.guardrails.tool_name_extraction import extract_request_tool_names + + cases = [ + ("OpenAI chat tools", "/v1/chat/completions", {"tools": [{"type": "function", "function": {"name": "get_weather"}}]}), + ("OpenAI chat functions", "/v1/chat/completions", {"functions": [{"name": "run_sql"}]}), + ("OpenAI responses function", "/v1/responses", {"tools": [{"type": "function", "name": "get_current_weather"}]}), + ("OpenAI responses MCP", "/v1/responses", {"tools": [{"type": "mcp", "server_label": "dmcp"}]}), + ("Anthropic", "/v1/messages", {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]}), + ("Google generateContent", "/generate_content", {"tools": [{"functionDeclarations": [{"name": "schedule_meeting"}]}]}), + ("MCP call_tool", "/mcp/call_tool", {"name": "my_tool", "arguments": {}}), + ("Non-tool route", "/v1/embeddings", {"tools": [{"type": "function", "function": {"name": "x"}}]}), + ] + print("=== extract_request_tool_names(route, data) ===\n") + for label, route, data in cases: + names = extract_request_tool_names(route, data) + print(f" {label}: {names}") + print() + + +async def test_check_tools_allowlist(): + """Test check_tools_allowlist with mock tokens.""" + from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import check_tools_allowlist + + def token(metadata=None, team_metadata=None): + return UserAPIKeyAuth( + api_key="test-key", + user_id="user", + team_id="team", + org_id=None, + models=["*"], + metadata=metadata or {}, + team_metadata=team_metadata or {}, + ) + + print("=== check_tools_allowlist (auth) ===\n") + + # No allowlist -> pass + await check_tools_allowlist( + request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]}, + valid_token=token(), + team_object=None, + route="/v1/chat/completions", + ) + print(" No allowlist, body has tools: PASS") + + # Allowed tool -> pass + await check_tools_allowlist( + request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]}, + valid_token=token(metadata={"allowed_tools": ["get_weather"]}), + team_object=None, + route="/v1/chat/completions", + ) + print(" allowed_tools=['get_weather'], body has get_weather: PASS") + + # Disallowed tool -> raise + try: + await check_tools_allowlist( + request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]}, + valid_token=token(metadata={"allowed_tools": ["other_tool"]}), + team_object=None, + route="/v1/chat/completions", + ) + print(" DISALLOWED: expected ProxyException") + except ProxyException as e: + if e.type == ProxyErrorTypes.tool_access_denied: + print(" allowed_tools=['other_tool'], body has get_weather: PASS (raised tool_access_denied)") + else: + print(f" Unexpected ProxyException type: {e.type}") + except Exception as e: + print(f" Unexpected: {e}") + + # Team allowlist when key empty + await check_tools_allowlist( + request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]}, + valid_token=token(team_metadata={"allowed_tools": ["get_weather"]}), + team_object=None, + route="/v1/chat/completions", + ) + print(" team_metadata.allowed_tools=['get_weather']: PASS") + print() + + +def main(): + print("Tool allowlist / tool name extraction – script checks\n") + test_extraction() + asyncio.run(test_check_tools_allowlist()) + print("Done. For full unit tests run:") + print(" poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v") + + +if __name__ == "__main__": + main() diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 38baaedef14..4f0459d5eca 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -903,3 +903,103 @@ async def test_anthropic_cache_control_hook_document_analysis_multiple_pages(): if isinstance(item, dict) and "cachePoint" in item ) assert cache_control_count == 1, f"Expected exactly 1 cache control point (last item only), found {cache_control_count}. Before fix, this would be 6 (one for each content item)." + + +def test_gemini_cache_control_injection_points_detected(): + """ + Test that cache_control_injection_points work for Gemini models. + + Verifies the full flow: + 1. The hook injects cache_control markers on string-content messages + 2. is_cached_message() detects the injected markers (message-level cache_control) + 3. separate_cached_messages() correctly separates the messages + + Fixes GitHub issue #18519. + """ + from litellm.llms.vertex_ai.context_caching.transformation import ( + separate_cached_messages, + ) + from litellm.utils import is_cached_message + + hook = AnthropicCacheControlHook() + + # Simulate messages as they would appear for a Gemini call with string content + messages: List[AllMessageValues] = [ + { + "role": "system", + "content": "You are a helpful assistant that analyzes legal documents.", + }, + { + "role": "user", + "content": "What are the key terms?", + }, + ] + + # Simulate what the hook does: inject cache_control on the system message + injection_points = [{"location": "message", "role": "system"}] + + # Manually apply the hook's logic for the system message (string content case) + # The hook sets message["cache_control"] = {"type": "ephemeral"} for string content + hook._safe_insert_cache_control_in_message( + message=messages[0], + control={"type": "ephemeral"}, + ) + + # Verify the hook injected message-level cache_control (string content path) + assert messages[0].get("cache_control") == {"type": "ephemeral"} + + # Verify is_cached_message detects message-level cache_control + assert is_cached_message(messages[0]) is True + assert is_cached_message(messages[1]) is False + + # Verify separate_cached_messages correctly separates them + cached, non_cached = separate_cached_messages(messages) + assert len(cached) == 1 + assert cached[0]["role"] == "system" + assert len(non_cached) == 1 + assert non_cached[0]["role"] == "user" + + +def test_gemini_cache_control_injection_list_content_detected(): + """ + Test that cache_control_injection_points work for Gemini models + when the message content is a list (not string). + """ + from litellm.llms.vertex_ai.context_caching.transformation import ( + separate_cached_messages, + ) + from litellm.utils import is_cached_message + + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Analyze legal documents carefully."}, + ], + }, + { + "role": "user", + "content": "What are the key terms?", + }, + ] + + # Apply the hook's logic for list content - sets cache_control on last item + hook._safe_insert_cache_control_in_message( + message=messages[0], + control={"type": "ephemeral"}, + ) + + # Verify cache_control was set on the last content item + assert messages[0]["content"][-1]["cache_control"] == {"type": "ephemeral"} + + # Verify is_cached_message detects content-item-level cache_control + assert is_cached_message(messages[0]) is True + assert is_cached_message(messages[1]) is False + + # Verify separate_cached_messages correctly separates them + cached, non_cached = separate_cached_messages(messages) + assert len(cached) == 1 + assert len(non_cached) == 1 diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index a47d026c169..3f8cbf12361 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -29,6 +29,15 @@ class TestContextCachingEndpoints: self.mock_client = MagicMock(spec=HTTPHandler) self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) + # Mock is_prompt_caching_valid_prompt to return True by default. + # This avoids token counting in unit tests. The min-token guard is + # tested explicitly in test_check_and_create_cache_skips_when_below_min_tokens. + self._token_check_patcher = patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt", + return_value=True, + ) + self._token_check_patcher.start() + # Sample messages for testing self.sample_messages = [ { @@ -56,6 +65,10 @@ class TestContextCachingEndpoints: self.sample_optional_params = {"tools": self.sample_tools.copy()} + def teardown_method(self): + """Teardown for each test method""" + self._token_check_patcher.stop() + @pytest.mark.parametrize( "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] ) @@ -787,6 +800,112 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + def test_check_and_create_cache_skips_when_below_min_tokens( + self, mock_separate, custom_llm_provider + ): + """Test that context caching is skipped when cached content is below 1024 tokens. + + Gemini requires a minimum of 1024 tokens for context caching. If the cached + content is too small, the request should proceed without caching instead of + failing with a Gemini API error. + """ + # Stop the default mock so the real token count check runs + self._token_check_patcher.stop() + + short_cached_messages = [ + { + "role": "system", + "content": "You are a helpful assistant.", + "cache_control": {"type": "ephemeral"}, + } + ] + non_cached_messages = [ + {"role": "user", "content": "Hello"}, + ] + all_messages = short_cached_messages + non_cached_messages + mock_separate.return_value = (short_cached_messages, non_cached_messages) + optional_params = self.sample_optional_params.copy() + + result = self.context_caching.check_and_create_cache( + messages=all_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + cached_content=None, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="test_token", + ) + + messages, returned_params, returned_cache = result + assert messages == all_messages + assert returned_cache is None + + # Restart the patcher so teardown_method can stop it cleanly + self._token_check_patcher.start() + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + @pytest.mark.asyncio + async def test_async_check_and_create_cache_skips_when_below_min_tokens( + self, mock_separate, custom_llm_provider + ): + """Test that async context caching is skipped when cached content is below 1024 tokens.""" + # Stop the default mock so the real token count check runs + self._token_check_patcher.stop() + + short_cached_messages = [ + { + "role": "system", + "content": "You are a helpful assistant.", + "cache_control": {"type": "ephemeral"}, + } + ] + non_cached_messages = [ + {"role": "user", "content": "Hello"}, + ] + all_messages = short_cached_messages + non_cached_messages + mock_separate.return_value = (short_cached_messages, non_cached_messages) + optional_params = self.sample_optional_params.copy() + + result = await self.context_caching.async_check_and_create_cache( + messages=all_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + cached_content=None, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="test_token", + ) + + messages, returned_params, returned_cache = result + assert messages == all_messages + assert returned_cache is None + + # Restart the patcher so teardown_method can stop it cleanly + self._token_check_patcher.start() + class TestCheckCachePagination: """Test pagination logic in check_cache and async_check_cache methods.""" diff --git a/tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py b/tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py new file mode 100644 index 00000000000..68d5e2035f7 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_file_retrieve_provider_routing.py @@ -0,0 +1,127 @@ +""" +Tests for Fix 1: file_retrieve Literal type was missing 'vertex_ai' and 'gemini', +causing a type mismatch when afile_retrieve delegated to the sync function. +""" + +import pytest +from unittest.mock import MagicMock, patch + +from litellm.files.main import file_retrieve + + +class TestFileRetrieveProviderRouting: + """ + Verify that file_retrieve accepts 'vertex_ai' and 'gemini' providers and + routes them through ProviderConfigManager / base_llm_http_handler. + """ + + def _make_mock_file_object(self): + mock = MagicMock() + mock.model_dump.return_value = { + "id": "gs://my-bucket/file.jsonl", + "object": "file", + "bytes": 1024, + "created_at": 0, + "filename": "file.jsonl", + "purpose": "batch", + "status": "processed", + } + return mock + + def test_should_route_vertex_ai_through_provider_config(self): + """ + Regression: file_retrieve Literal type was missing 'vertex_ai', + so passing custom_llm_provider='vertex_ai' would fail type-checking + and potentially cause a routing failure at runtime. + """ + mock_file = self._make_mock_file_object() + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ) as mock_retrieve: + result = file_retrieve( + file_id="gs://my-bucket/file.jsonl", + custom_llm_provider="vertex_ai", + ) + + mock_retrieve.assert_called_once() + assert result is not None + + def test_should_route_gemini_through_provider_config(self): + """ + Regression: file_retrieve Literal type was also missing 'gemini'. + """ + mock_file = self._make_mock_file_object() + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ) as mock_retrieve: + result = file_retrieve( + file_id="some-gemini-file-id", + custom_llm_provider="gemini", + ) + + mock_retrieve.assert_called_once() + assert result is not None + + def test_should_pass_file_id_to_handler_for_vertex_ai(self): + """Verify the file_id is forwarded correctly to the underlying handler.""" + mock_file = self._make_mock_file_object() + expected_file_id = "gs://my-bucket/path/to/file.jsonl" + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ) as mock_retrieve: + file_retrieve( + file_id=expected_file_id, + custom_llm_provider="vertex_ai", + ) + + call_kwargs = mock_retrieve.call_args.kwargs + assert call_kwargs.get("file_id") == expected_file_id + + def test_should_not_raise_bad_request_for_vertex_ai(self): + """ + Before the fix, vertex_ai fell through to the else-branch which raised + BadRequestError. Verify it no longer does. + """ + import litellm + + mock_file = self._make_mock_file_object() + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ): + try: + file_retrieve( + file_id="gs://my-bucket/file.jsonl", + custom_llm_provider="vertex_ai", + ) + except litellm.exceptions.BadRequestError as e: + pytest.fail( + f"file_retrieve raised BadRequestError for vertex_ai: {e}" + ) + + def test_should_not_raise_bad_request_for_gemini(self): + """Same as above but for 'gemini'.""" + import litellm + + mock_file = self._make_mock_file_object() + + with patch( + "litellm.files.main.base_llm_http_handler.retrieve_file", + return_value=mock_file, + ): + try: + file_retrieve( + file_id="some-file-id", + custom_llm_provider="gemini", + ) + except litellm.exceptions.BadRequestError as e: + pytest.fail( + f"file_retrieve raised BadRequestError for gemini: {e}" + ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 6f1d753484d..598ad255aca 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -167,7 +167,7 @@ class TestTransformDeleteFile: assert isinstance(result, FileDeleted) assert result.deleted is True assert result.object == "file" - assert "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" in result.id + assert result.id == "gs://my-bucket/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" def test_should_fallback_to_deleted_id_when_no_request(self, config): raw_response = MagicMock(spec=httpx.Response) @@ -182,3 +182,49 @@ class TestTransformDeleteFile: assert isinstance(result, FileDeleted) assert result.id == "deleted" assert result.deleted is True + + def test_should_include_bucket_name_in_reconstructed_delete_id(self, config): + """ + Regression: the old code split on /o/ only, dropping the bucket from + the reconstructed gs:// URI. e.g. gs://path/to/file instead of + gs://my-bucket/path/to/file. + """ + raw_response = MagicMock(spec=httpx.Response) + mock_request = MagicMock() + encoded_object = urllib.parse.quote("path/to/file.jsonl", safe="") + mock_request.url = ( + f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_object}" + ) + raw_response.request = mock_request + + result = config.transform_delete_file_response( + raw_response=raw_response, + logging_obj=MagicMock(), + litellm_params={}, + ) + + assert result.id == "gs://my-bucket/path/to/file.jsonl" + + def test_should_include_bucket_in_nested_object_path(self, config): + """Verify bucket extraction works with deeply nested GCS object paths.""" + raw_response = MagicMock(spec=httpx.Response) + mock_request = MagicMock() + encoded_object = urllib.parse.quote( + "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123", + safe="", + ) + mock_request.url = ( + f"https://storage.googleapis.com/storage/v1/b/prod-bucket/o/{encoded_object}" + ) + raw_response.request = mock_request + + result = config.transform_delete_file_response( + raw_response=raw_response, + logging_obj=MagicMock(), + litellm_params={}, + ) + + assert result.id == ( + "gs://prod-bucket/litellm-vertex-files/publishers/google/" + "models/gemini-2.0-flash-001/abc-123" + ) diff --git a/tests/test_litellm/proxy/db/test_tool_registry_writer.py b/tests/test_litellm/proxy/db/test_tool_registry_writer.py index 44f9e32058a..1b1ee7afcba 100644 --- a/tests/test_litellm/proxy/db/test_tool_registry_writer.py +++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py @@ -1,6 +1,6 @@ """ Unit tests for tool_registry_writer.py — uses a mock prisma client -that exposes execute_raw / query_raw (matching the actual raw-SQL implementation). +that exposes litellm_tooltable.upsert / find_many / find_unique. """ import os @@ -13,21 +13,28 @@ import pytest sys.path.insert(0, os.path.abspath("../../..")) from litellm.proxy.db.tool_registry_writer import ( + ToolPolicyRegistry, batch_upsert_tools, get_tool, + get_tool_policy_registry, get_tools_by_names, list_tools, update_tool_policy, ) -def _make_prisma(query_rows=None): - """Return a minimal mock prisma_client with execute_raw / query_raw.""" - default_row = { +def _mock_row(**kwargs): + """Build a row-like object with real attributes (no MagicMock) for _row_to_model.""" + + class Row: + pass + + default = { "tool_id": "uuid-1", "tool_name": "my_tool", "origin": "user_defined", - "call_policy": "untrusted", + "input_policy": "untrusted", + "output_policy": "untrusted", "call_count": 1, "assignments": {}, "key_hash": None, @@ -38,31 +45,54 @@ def _make_prisma(query_rows=None): "created_by": None, "updated_by": None, } - rows = query_rows if query_rows is not None else [default_row] + default.update(kwargs) + row = Row() + for k, v in default.items(): + setattr(row, k, v) + return row + +def _make_prisma( + *, + upsert_return=None, + find_many_rows=None, + find_unique_row=None, +): + """Return a mock prisma_client with litellm_tooltable.upsert, find_many, find_unique.""" prisma = MagicMock() - prisma.db.execute_raw = AsyncMock(return_value=None) - prisma.db.query_raw = AsyncMock(return_value=rows) + prisma.db.litellm_tooltable = MagicMock() + prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) + prisma.db.litellm_tooltable.find_many = AsyncMock( + return_value=find_many_rows if find_many_rows is not None else [] + ) + prisma.db.litellm_tooltable.find_unique = AsyncMock( + return_value=find_unique_row + ) return prisma @pytest.mark.asyncio -async def test_batch_upsert_tools_calls_execute_raw(): +async def test_batch_upsert_tools_calls_upsert(): prisma = _make_prisma() items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] await batch_upsert_tools(prisma, items) - prisma.db.execute_raw.assert_awaited_once() - call_args = prisma.db.execute_raw.call_args - sql = call_args.args[0] - assert "LiteLLM_ToolTable" in sql - assert "ON CONFLICT" in sql + prisma.db.litellm_tooltable.upsert.assert_awaited_once() + call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs + assert call_kw["where"] == {"tool_name": "tool_a"} + assert call_kw["data"]["create"]["tool_name"] == "tool_a" + assert call_kw["data"]["create"]["origin"] == "mcp_server" + assert call_kw["data"]["create"]["input_policy"] == "untrusted" + assert call_kw["data"]["create"]["output_policy"] == "untrusted" + assert call_kw["data"]["create"]["call_count"] == 1 + assert call_kw["data"]["update"]["call_count"] == {"increment": 1} + assert "updated_at" in call_kw["data"]["update"] @pytest.mark.asyncio async def test_batch_upsert_tools_empty_list(): prisma = _make_prisma() await batch_upsert_tools(prisma, []) - prisma.db.execute_raw.assert_not_awaited() + prisma.db.litellm_tooltable.upsert.assert_not_awaited() @pytest.mark.asyncio @@ -70,123 +100,120 @@ async def test_batch_upsert_tools_skips_empty_names(): prisma = _make_prisma() items = [{"tool_name": "", "origin": None}, {"tool_name": None}] # type: ignore[list-item] await batch_upsert_tools(prisma, items) - prisma.db.execute_raw.assert_not_awaited() + prisma.db.litellm_tooltable.upsert.assert_not_awaited() @pytest.mark.asyncio -async def test_batch_upsert_multiple_tools_calls_execute_raw_per_tool(): +async def test_batch_upsert_multiple_tools_calls_upsert_per_tool(): prisma = _make_prisma() items = [ {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}, {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"}, ] await batch_upsert_tools(prisma, items) - assert prisma.db.execute_raw.await_count == 2 + assert prisma.db.litellm_tooltable.upsert.await_count == 2 + calls = prisma.db.litellm_tooltable.upsert.call_args_list + assert calls[0].kwargs["where"]["tool_name"] == "tool_a" + assert calls[1].kwargs["where"]["tool_name"] == "tool_b" @pytest.mark.asyncio async def test_list_tools_no_filter(): - row = { - "tool_id": "id1", - "tool_name": "tool_a", - "origin": "mcp", - "call_policy": "untrusted", - "call_count": 5, - "assignments": {}, - "key_hash": None, - "team_id": None, - "key_alias": None, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - "created_by": None, - "updated_by": None, - } - prisma = _make_prisma(query_rows=[row]) + row = _mock_row( + tool_id="id1", + tool_name="tool_a", + origin="mcp", + input_policy="untrusted", + output_policy="untrusted", + call_count=5, + ) + prisma = _make_prisma(find_many_rows=[row]) result = await list_tools(prisma) assert len(result) == 1 assert result[0].tool_name == "tool_a" assert result[0].call_count == 5 - prisma.db.query_raw.assert_awaited_once() + prisma.db.litellm_tooltable.find_many.assert_awaited_once() + call_kw = prisma.db.litellm_tooltable.find_many.call_args.kwargs + assert call_kw["where"] == {} + assert call_kw["order"] == {"created_at": "desc"} @pytest.mark.asyncio -async def test_list_tools_with_policy_filter(): - row = { - "tool_id": "id1", - "tool_name": "blocked_tool", - "origin": None, - "call_policy": "blocked", - "call_count": 2, - "assignments": None, - "key_hash": None, - "team_id": None, - "key_alias": None, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - "created_by": None, - "updated_by": None, - } - prisma = _make_prisma(query_rows=[row]) - result = await list_tools(prisma, call_policy="blocked") - assert result[0].call_policy == "blocked" - call_args = prisma.db.query_raw.call_args - sql = call_args.args[0] - assert "WHERE call_policy" in sql +async def test_list_tools_with_input_policy_filter(): + row = _mock_row( + tool_id="id1", + tool_name="blocked_tool", + origin=None, + input_policy="blocked", + output_policy="untrusted", + call_count=2, + assignments=None, + ) + prisma = _make_prisma(find_many_rows=[row]) + result = await list_tools(prisma, input_policy="blocked") + assert result[0].input_policy == "blocked" + call_kw = prisma.db.litellm_tooltable.find_many.call_args.kwargs + assert call_kw["where"] == {"input_policy": "blocked"} @pytest.mark.asyncio async def test_get_tool_found(): - prisma = _make_prisma() + row = _mock_row(tool_name="my_tool") + prisma = _make_prisma(find_unique_row=row) result = await get_tool(prisma, "my_tool") assert result is not None assert result.tool_name == "my_tool" - prisma.db.query_raw.assert_awaited_once() + prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with( + where={"tool_name": "my_tool"} + ) @pytest.mark.asyncio async def test_get_tool_not_found(): - prisma = _make_prisma(query_rows=[]) + prisma = _make_prisma(find_unique_row=None) result = await get_tool(prisma, "nonexistent") assert result is None @pytest.mark.asyncio -async def test_update_tool_policy_calls_execute_raw(): - row = { - "tool_id": "uuid-1", - "tool_name": "my_tool", - "origin": "user_defined", - "call_policy": "blocked", - "call_count": 1, - "assignments": {}, - "key_hash": None, - "team_id": None, - "key_alias": None, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - "created_by": None, - "updated_by": "admin", - } - prisma = _make_prisma(query_rows=[row]) - result = await update_tool_policy(prisma, "my_tool", "blocked", "admin") +async def test_update_tool_policy_calls_upsert_then_get_tool(): + row = _mock_row( + tool_name="my_tool", + input_policy="blocked", + output_policy="untrusted", + updated_by="admin", + ) + prisma = _make_prisma(find_unique_row=row) + result = await update_tool_policy( + prisma, "my_tool", updated_by="admin", input_policy="blocked" + ) assert result is not None - assert result.call_policy == "blocked" - prisma.db.execute_raw.assert_awaited_once() - call_args = prisma.db.execute_raw.call_args - sql = call_args.args[0] - assert "ON CONFLICT" in sql - assert "call_policy" in sql + assert result.input_policy == "blocked" + prisma.db.litellm_tooltable.upsert.assert_awaited_once() + call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs + assert call_kw["where"] == {"tool_name": "my_tool"} + assert call_kw["data"]["update"]["input_policy"] == "blocked" + assert call_kw["data"]["update"]["updated_by"] == "admin" + prisma.db.litellm_tooltable.find_unique.assert_awaited_with( + where={"tool_name": "my_tool"} + ) @pytest.mark.asyncio async def test_get_tools_by_names_returns_policy_map(): rows = [ - {"tool_name": "tool_a", "call_policy": "trusted"}, - {"tool_name": "tool_b", "call_policy": "blocked"}, + _mock_row(tool_name="tool_a", input_policy="trusted", output_policy="untrusted"), + _mock_row(tool_name="tool_b", input_policy="blocked", output_policy="untrusted"), ] - prisma = _make_prisma(query_rows=rows) + prisma = _make_prisma(find_many_rows=rows) result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) - assert result == {"tool_a": "trusted", "tool_b": "blocked"} + assert result == { + "tool_a": ("trusted", "untrusted"), + "tool_b": ("blocked", "untrusted"), + } + prisma.db.litellm_tooltable.find_many.assert_awaited_once_with( + where={"tool_name": {"in": ["tool_a", "tool_b"]}} + ) @pytest.mark.asyncio @@ -194,4 +221,71 @@ async def test_get_tools_by_names_empty_list(): prisma = _make_prisma() result = await get_tools_by_names(prisma, []) assert result == {} - prisma.db.query_raw.assert_not_awaited() + prisma.db.litellm_tooltable.find_many.assert_not_awaited() + + +# --- ToolPolicyRegistry --- + + +def _mock_tool_row( + tool_name: str, + input_policy: str = "untrusted", + output_policy: str = "untrusted", +): + row = MagicMock() + row.tool_name = tool_name + row.input_policy = input_policy + row.output_policy = output_policy + return row + + +def _mock_perm_row(object_permission_id: str, blocked_tools: list): + row = MagicMock() + row.object_permission_id = object_permission_id + row.blocked_tools = blocked_tools + return row + + +@pytest.mark.asyncio +async def test_tool_policy_registry_sync_and_get_effective_policies(): + """Registry syncs from DB; get_effective_policies returns merged blocked + global.""" + prisma = MagicMock() + prisma.db.litellm_tooltable.find_many = AsyncMock( + return_value=[ + _mock_tool_row("tool_a", input_policy="trusted"), + _mock_tool_row("tool_b", input_policy="blocked"), + _mock_tool_row("tool_c", input_policy="untrusted"), + ] + ) + prisma.db.litellm_objectpermissiontable.find_many = AsyncMock( + return_value=[ + _mock_perm_row("op-key-1", ["tool_a"]), + _mock_perm_row("op-team-1", ["tool_c"]), + ] + ) + registry = get_tool_policy_registry() + await registry.sync_tool_policy_from_db(prisma) + assert registry.is_initialized() + # Key blocked: tool_a. Team blocked: tool_c. Global: tool_b blocked. + result = registry.get_effective_policies( + ["tool_a", "tool_b", "tool_c"], + object_permission_id="op-key-1", + team_object_permission_id="op-team-1", + ) + assert result["tool_a"] == "blocked" + assert result["tool_b"] == "blocked" + assert result["tool_c"] == "blocked" + # No op ids: only global + result_global = registry.get_effective_policies(["tool_a", "tool_b", "tool_c"]) + assert result_global["tool_a"] == "trusted" + assert result_global["tool_b"] == "blocked" + assert result_global["tool_c"] == "untrusted" + + +@pytest.mark.asyncio +async def test_tool_policy_registry_not_initialized_returns_untrusted(): + """When not synced, get_effective_policies still returns untrusted for unknown tools.""" + registry = ToolPolicyRegistry() + assert not registry.is_initialized() + result = registry.get_effective_policies(["unknown_tool"]) + assert result == {"unknown_tool": "untrusted"} diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py index c6a81efbf0b..943a8d4be75 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py @@ -12,9 +12,8 @@ from fastapi import HTTPException sys.path.insert(0, os.path.abspath("../../../../../..")) -from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( - ToolPolicyGuardrail, -) +from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import \ + ToolPolicyGuardrail from litellm.types.guardrails import GuardrailEventHooks @@ -70,10 +69,21 @@ async def test_no_tool_calls_in_response_passes_through(guardrail): assert result is inputs +def _registry_mock(policy_map: dict): + """Return a mock registry with is_initialized=True and get_effective_policies returning policy_map.""" + reg = MagicMock() + reg.is_initialized.return_value = True + reg.get_effective_policies.return_value = policy_map + return reg + + @pytest.mark.asyncio async def test_untrusted_tools_pass_through(guardrail): policy_map = {"search": "untrusted", "read_file": "trusted"} - with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock(policy_map), + ): inputs: Any = _tool_request_inputs(["search", "read_file"]) result = await guardrail.apply_guardrail( inputs=inputs, request_data={}, input_type="request" @@ -84,7 +94,10 @@ async def test_untrusted_tools_pass_through(guardrail): @pytest.mark.asyncio async def test_blocked_tool_in_request_raises_http_exception(guardrail): policy_map = {"dangerous_tool": "blocked"} - with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock(policy_map), + ): inputs: Any = _tool_request_inputs(["dangerous_tool"]) with pytest.raises(HTTPException) as exc_info: await guardrail.apply_guardrail( @@ -97,7 +110,10 @@ async def test_blocked_tool_in_request_raises_http_exception(guardrail): @pytest.mark.asyncio async def test_blocked_tool_in_response_raises_http_exception(guardrail): policy_map = {"exfil_tool": "blocked"} - with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock(policy_map), + ): inputs: Any = _tool_response_inputs(["exfil_tool"]) with pytest.raises(HTTPException) as exc_info: await guardrail.apply_guardrail( @@ -110,7 +126,10 @@ async def test_blocked_tool_in_response_raises_http_exception(guardrail): @pytest.mark.asyncio async def test_mixed_blocked_and_allowed_raises_for_blocked(guardrail): policy_map = {"safe_tool": "trusted", "bad_tool": "blocked"} - with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock(policy_map), + ): inputs: Any = _tool_request_inputs(["safe_tool", "bad_tool"]) with pytest.raises(HTTPException) as exc_info: await guardrail.apply_guardrail( @@ -123,8 +142,11 @@ async def test_mixed_blocked_and_allowed_raises_for_blocked(guardrail): @pytest.mark.asyncio async def test_tool_not_in_db_passes_through(guardrail): - """Tools not found in the DB (no entry) should not be blocked.""" - with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value={})): + """When registry returns no policy (or empty), tools are not blocked.""" + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock({}), + ): inputs: Any = _tool_request_inputs(["unknown_tool"]) result = await guardrail.apply_guardrail( inputs=inputs, request_data={}, input_type="request" @@ -133,43 +155,30 @@ async def test_tool_not_in_db_passes_through(guardrail): @pytest.mark.asyncio -async def test_get_policies_cached_uses_cache(guardrail): - """Second call with same tool names should return the cached result.""" - policy_map = {"tool_a": "trusted"} +async def test_registry_not_initialized_passes_through(guardrail): + """When registry is not initialized, no tools are blocked (empty policy map).""" + reg = MagicMock() + reg.is_initialized.return_value = False with patch( - "litellm.proxy.db.tool_registry_writer.get_tools_by_names", - new=AsyncMock(return_value=policy_map), - ) as mock_db, patch( - "litellm.proxy.proxy_server.prisma_client", - new=MagicMock(), + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=reg, ): - # first call — should hit DB - result1 = await guardrail._get_policies_cached(["tool_a"]) - assert result1 == policy_map - - # second call — should hit cache, not DB again - result2 = await guardrail._get_policies_cached(["tool_a"]) - assert result2 == policy_map - - assert mock_db.call_count == 1 - - -@pytest.mark.asyncio -async def test_get_policies_cached_no_prisma(guardrail): - """Without a prisma client, returns empty dict.""" - with patch( - "litellm.proxy.proxy_server.prisma_client", - None, - ): - result = await guardrail._get_policies_cached(["tool_a"]) - assert result == {} + inputs: Any = _tool_request_inputs(["any_tool"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + reg.get_effective_policies.assert_not_called() @pytest.mark.asyncio async def test_response_tool_calls_as_objects(guardrail): """tool_calls that are objects (not dicts) with .function.name should work.""" policy_map = {"obj_tool": "blocked"} - with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock(policy_map), + ): fn = MagicMock() fn.name = "obj_tool" tc = MagicMock() diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 62a6e777b0d..ca224726361 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1112,6 +1112,47 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker): assert result.guardrail_definition_location == "db" +class TestBuildFieldDict: + """Test _build_field_dict handles both enum and string ui_type values.""" + + def test_build_field_dict_with_string_ui_type(self): + """Test that _build_field_dict works when ui_type is a plain string (e.g. BlockCodeExecutionGuardrailConfigModel).""" + from unittest.mock import MagicMock + + from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict + + field = MagicMock() + field.json_schema_extra = {"ui_type": "multiselect", "options": ["python", "javascript"]} + + result = _build_field_dict( + field=field, + field_annotation=str, + description="Test field", + required=False, + ) + + assert result["type"] == "multiselect" + assert result["description"] == "Test field" + + def test_build_field_dict_with_enum_ui_type(self): + """Test that _build_field_dict works when ui_type is a GuardrailParamUITypes enum.""" + from unittest.mock import MagicMock + + from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict + from litellm.types.guardrails import GuardrailParamUITypes + + field = MagicMock() + field.json_schema_extra = {"ui_type": GuardrailParamUITypes.BOOL} + + result = _build_field_dict( + field=field, + field_annotation=bool, + description="Test bool field", + required=True, + ) + + assert result["type"] == "bool" + assert result["required"] is True # --- Team guardrail registration (register / submissions) --- MOCK_REGISTER_REQUEST = RegisterGuardrailRequest( @@ -1571,4 +1612,4 @@ async def test_list_submissions_summary_counts_unaffected_by_filters(mocker): assert len(result.submissions) == 1 # filtered assert result.summary.total == 2 # unfiltered assert result.summary.pending_review == 1 - assert result.summary.active == 1 \ No newline at end of file + assert result.summary.active == 1 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 9b905d24fd1..ba1084eafe0 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,7 +1,7 @@ import copy import datetime from typing import AsyncGenerator -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import Request, status @@ -13,13 +13,13 @@ from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ProxyConfig, - _add_dd_apm_tags_for_litellm_call_id, _extract_error_from_sse_chunk, _get_cost_breakdown_from_logging_obj, _override_openai_response_model, _parse_event_data_for_error, create_response, ) +from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.utils import ProxyLogging @@ -82,13 +82,15 @@ class TestProxyBaseLLMRequestProcessing: def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch): mock_set_active_span_tag = MagicMock(return_value=True) + import litellm.proxy.dd_span_tagger + monkeypatch.setattr( - litellm.proxy.common_request_processing, + litellm.proxy.dd_span_tagger, "set_active_span_tag", mock_set_active_span_tag, ) - _add_dd_apm_tags_for_litellm_call_id("test-call-id") + DDSpanTagger.tag_call_id("test-call-id") mock_set_active_span_tag.assert_called_once_with( "litellm.call_id", "test-call-id" @@ -1564,3 +1566,59 @@ class TestStreamingOverheadHeader: "It was missing — this is the streaming overhead header regression." ) assert custom_headers["x-litellm-overhead-duration-ms"] == "55.3" + + +class TestDDSpanTaggerTagRequest: + """Tests for DDSpanTagger.tag_request - key/model DD span tagging.""" + + def _make_user_api_key_dict(self, key_alias=None, token=None): + from litellm.proxy._types import UserAPIKeyAuth + + d = UserAPIKeyAuth() + d.key_alias = key_alias + d.token = token + return d + + def test_tags_key_alias_and_model(self): + """key_alias and requested_model are set on the span when present.""" + user_key = self._make_user_api_key_dict(key_alias="my-prod-key", token="hashed123") + + with patch( + "litellm.proxy.dd_span_tagger.set_active_span_tag" + ) as mock_set_tag: + DDSpanTagger.tag_request( + user_api_key_dict=user_key, + requested_model="gpt-4o", + ) + + mock_set_tag.assert_any_call("litellm.key_alias", "my-prod-key") + mock_set_tag.assert_any_call("litellm.key_hash", "hashed123") + mock_set_tag.assert_any_call("litellm.requested_model", "gpt-4o") + + def test_no_tags_when_key_absent(self): + """No key tags are set when key_alias and token are None (e.g. 401 path).""" + user_key = self._make_user_api_key_dict(key_alias=None, token=None) + + with patch( + "litellm.proxy.dd_span_tagger.set_active_span_tag" + ) as mock_set_tag: + DDSpanTagger.tag_request( + user_api_key_dict=user_key, + requested_model=None, + ) + + mock_set_tag.assert_not_called() + + def test_only_model_tagged_when_no_key_info(self): + """requested_model is tagged even when there's no key info.""" + user_key = self._make_user_api_key_dict(key_alias=None, token=None) + + with patch( + "litellm.proxy.dd_span_tagger.set_active_span_tag" + ) as mock_set_tag: + DDSpanTagger.tag_request( + user_api_key_dict=user_key, + requested_model="claude-3-5-sonnet", + ) + + mock_set_tag.assert_called_once_with("litellm.requested_model", "claude-3-5-sonnet") diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 8abc6bfe077..bc13cea939e 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -11,11 +11,16 @@ from fastapi import Request import litellm from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import ( - KeyAndTeamLoggingSettings, LiteLLMProxyRequestSetup, - _get_dynamic_logging_metadata, _get_enforced_params, - _get_metadata_variable_name, _update_model_if_key_alias_exists, - add_guardrails_from_policy_engine, add_litellm_data_to_request, - check_if_token_is_service_account) + KeyAndTeamLoggingSettings, + LiteLLMProxyRequestSetup, + _get_dynamic_logging_metadata, + _get_enforced_params, + _get_metadata_variable_name, + _update_model_if_key_alias_exists, + add_guardrails_from_policy_engine, + add_litellm_data_to_request, + check_if_token_is_service_account, +) sys.path.insert( 0, os.path.abspath("../../..") @@ -154,8 +159,7 @@ def test_get_enforced_params( @pytest.mark.asyncio async def test_add_litellm_data_to_request_parses_string_metadata(): - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Setup request_mock = MagicMock(spec=Request) @@ -201,8 +205,7 @@ async def test_add_litellm_data_to_request_parses_string_metadata(): @pytest.mark.asyncio async def test_add_litellm_data_to_request_user_spend_and_budget(): - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request request_mock = MagicMock(spec=Request) request_mock.url.path = "/v1/completions" @@ -240,8 +243,7 @@ async def test_add_litellm_data_to_request_user_spend_and_budget(): @pytest.mark.asyncio async def test_add_litellm_data_to_request_audio_transcription_multipart(): - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Setup request mock for /v1/audio/transcriptions request_mock = MagicMock(spec=Request) @@ -306,8 +308,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks(): """ Test that litellm_disabled_callbacks from key metadata is properly added to the request data. """ - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Setup mock request request_mock = MagicMock(spec=Request) @@ -360,8 +361,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks_empty(): """ Test that litellm_disabled_callbacks is not added when it's empty. """ - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Setup mock request request_mock = MagicMock(spec=Request) @@ -413,8 +413,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks_not_present(): """ Test that litellm_disabled_callbacks is not added when it's not present in metadata. """ - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Setup mock request request_mock = MagicMock(spec=Request) @@ -466,8 +465,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks_invalid_type(): """ Test that litellm_disabled_callbacks is not added when it's not a list. """ - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Setup mock request request_mock = MagicMock(spec=Request) @@ -519,8 +517,7 @@ async def test_add_litellm_data_to_request_disabled_callbacks_with_logging_setti """ Test that litellm_disabled_callbacks works correctly alongside logging settings. """ - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Setup mock request request_mock = MagicMock(spec=Request) @@ -1030,8 +1027,7 @@ from unittest.mock import AsyncMock from fastapi.responses import Response from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy.common_request_processing import \ - ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import ProxyLogging from litellm.types.utils import StandardLoggingPayload @@ -1149,6 +1145,47 @@ async def test_add_litellm_metadata_from_request_headers(): litellm.callbacks = original_callbacks +def test_add_litellm_metadata_from_request_headers_x_litellm_trace_id_sets_chain_id(): + """x-litellm-trace-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining.""" + headers = {"x-litellm-trace-id": "foo"} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["trace_id"] == "foo" + assert data["metadata"]["session_id"] == "foo" + assert data["litellm_session_id"] == "foo" + assert data["litellm_trace_id"] == "foo" + + +def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id(): + """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining.""" + headers = {"x-litellm-session-id": "bar"} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["trace_id"] == "bar" + assert data["metadata"]["session_id"] == "bar" + assert data["litellm_session_id"] == "bar" + assert data["litellm_trace_id"] == "bar" + + +def test_add_litellm_metadata_from_request_headers_both_headers_trace_id_precedence(): + """When both x-litellm-trace-id and x-litellm-session-id are present, trace-id takes precedence for chain_id.""" + headers = { + "x-litellm-trace-id": "trace-value", + "x-litellm-session-id": "session-value", + } + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["trace_id"] == "trace-value" + assert data["metadata"]["session_id"] == "trace-value" + assert data["litellm_session_id"] == "trace-value" + assert data["litellm_trace_id"] == "trace-value" + def test_get_internal_user_header_from_mapping_returns_expected_header(): mappings = [ @@ -1407,8 +1444,7 @@ async def test_embedding_header_forwarding_with_model_group(): importlib.reload(pre_call_utils_module) # Re-import the function after reload to get the fresh version - from litellm.proxy.litellm_pre_call_utils import \ - add_litellm_data_to_request + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request # Setup mock request for embeddings request_mock = MagicMock(spec=Request) @@ -1542,11 +1578,13 @@ async def test_add_guardrails_from_policy_engine(): Test that add_guardrails_from_policy_engine adds guardrails from matching policies and tracks applied policies in metadata. """ - from litellm.proxy.policy_engine.attachment_registry import \ - get_attachment_registry + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry - from litellm.types.proxy.policy_engine import (Policy, PolicyAttachment, - PolicyGuardrails) + from litellm.types.proxy.policy_engine import ( + Policy, + PolicyAttachment, + PolicyGuardrails, + ) # Setup test data data = { @@ -1659,8 +1697,7 @@ async def test_add_guardrails_from_policy_engine_policy_version_by_id(): Test that add_guardrails_from_policy_engine executes a specific policy version when policy_ is passed in the request body. """ - from litellm.proxy.policy_engine.attachment_registry import \ - get_attachment_registry + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails @@ -1729,6 +1766,7 @@ async def test_bearer_token_not_in_debug_logs(): """ import logging from io import StringIO + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ProxyConfig diff --git a/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py b/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py new file mode 100644 index 00000000000..4adc5acde8b --- /dev/null +++ b/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py @@ -0,0 +1,200 @@ +""" +Tests for tool allowlist enforcement (key/team metadata.allowed_tools). + +Covers: +- check_tools_allowlist: allowed, disallowed, no allowlist, non-tool routes +- extract_request_tool_names: OpenAI chat, responses, Anthropic, generate_content, MCP +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy._types import (ProxyErrorTypes, ProxyException, + UserAPIKeyAuth) +from litellm.proxy.auth.auth_checks import check_tools_allowlist +from litellm.proxy.guardrails.tool_name_extraction import ( + TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names) + + +def _token(metadata=None, team_metadata=None): + return UserAPIKeyAuth( + api_key="test-key", + user_id="user", + team_id="team", + org_id=None, + models=["*"], + metadata=metadata or {}, + team_metadata=team_metadata or {}, + ) + + +class TestExtractRequestToolNames: + """Test tool name extraction per API format.""" + + def test_openai_chat_tools(self): + data = { + "tools": [ + {"type": "function", "function": {"name": "get_weather"}}, + {"type": "function", "function": {"name": "run_sql"}}, + ] + } + assert extract_request_tool_names("/v1/chat/completions", data) == [ + "get_weather", + "run_sql", + ] + + def test_openai_chat_functions_legacy(self): + data = {"functions": [{"name": "get_weather"}, {"name": "run_sql"}]} + assert extract_request_tool_names("/v1/chat/completions", data) == [ + "get_weather", + "run_sql", + ] + + def test_openai_responses_function_tools(self): + data = { + "tools": [ + {"type": "function", "name": "get_current_weather", "description": "x"}, + ] + } + assert extract_request_tool_names("/v1/responses", data) == [ + "get_current_weather" + ] + + def test_openai_responses_mcp_tools(self): + data = { + "tools": [ + {"type": "mcp", "server_label": "dmcp", "server_url": "http://x"}, + ] + } + assert extract_request_tool_names("/v1/responses", data) == ["dmcp"] + + def test_anthropic_tools(self): + data = {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]} + assert extract_request_tool_names("/v1/messages", data) == [ + "get_weather", + "run_sql", + ] + + def test_generate_content_tools(self): + data = { + "tools": [ + { + "functionDeclarations": [ + {"name": "schedule_meeting", "description": "x"}, + ] + }, + ] + } + assert extract_request_tool_names("/generate_content", data) == [ + "schedule_meeting" + ] + + def test_mcp_call_tool_name(self): + data = {"name": "my_tool", "arguments": {}} + assert extract_request_tool_names("/mcp/call_tool", data) == ["my_tool"] + + def test_mcp_call_tool_mcp_tool_name(self): + data = {"mcp_tool_name": "other_tool"} + assert extract_request_tool_names("/mcp/call_tool", data) == ["other_tool"] + + def test_non_tool_route_returns_empty(self): + data = {"tools": [{"type": "function", "function": {"name": "x"}}]} + assert extract_request_tool_names("/v1/embeddings", data) == [] + + +class TestCheckToolsAllowlist: + """Test allowlist enforcement in auth (no DB in hot path).""" + + @pytest.mark.asyncio + async def test_no_allowlist_passes(self): + token = _token(metadata={}, team_metadata={}) + body = { + "tools": [{"type": "function", "function": {"name": "get_weather"}}] + } + await check_tools_allowlist( + request_body=body, + valid_token=token, + team_object=None, + route="/v1/chat/completions", + ) + + @pytest.mark.asyncio + async def test_allowed_tool_passes(self): + token = _token(metadata={"allowed_tools": ["get_weather"]}) + body = { + "tools": [{"type": "function", "function": {"name": "get_weather"}}] + } + await check_tools_allowlist( + request_body=body, + valid_token=token, + team_object=None, + route="/v1/chat/completions", + ) + + @pytest.mark.asyncio + async def test_disallowed_tool_raises(self): + token = _token(metadata={"allowed_tools": ["other_tool"]}) + body = { + "tools": [{"type": "function", "function": {"name": "get_weather"}}] + } + with pytest.raises(ProxyException) as exc_info: + await check_tools_allowlist( + request_body=body, + valid_token=token, + team_object=None, + route="/v1/chat/completions", + ) + assert exc_info.value.type == ProxyErrorTypes.tool_access_denied + assert "get_weather" in str(exc_info.value.message) + + @pytest.mark.asyncio + async def test_team_allowlist_used_when_key_empty(self): + token = _token( + metadata={}, + team_metadata={"allowed_tools": ["get_weather"]}, + ) + body = { + "tools": [{"type": "function", "function": {"name": "get_weather"}}] + } + await check_tools_allowlist( + request_body=body, + valid_token=token, + team_object=None, + route="/v1/chat/completions", + ) + + @pytest.mark.asyncio + async def test_key_allowlist_overrides_team(self): + token = _token( + metadata={"allowed_tools": ["get_weather"]}, + team_metadata={"allowed_tools": ["other_tool"]}, + ) + body = { + "tools": [{"type": "function", "function": {"name": "get_weather"}}] + } + await check_tools_allowlist( + request_body=body, + valid_token=token, + team_object=None, + route="/v1/chat/completions", + ) + + @pytest.mark.asyncio + async def test_valid_token_none_skips(self): + await check_tools_allowlist( + request_body={"tools": [{"type": "function", "function": {"name": "x"}}]}, + valid_token=None, + team_object=None, + route="/v1/chat/completions", + ) + + @pytest.mark.asyncio + async def test_no_tools_in_body_passes(self): + token = _token(metadata={"allowed_tools": ["get_weather"]}) + await check_tools_allowlist( + request_body={"messages": []}, + valid_token=token, + team_object=None, + route="/v1/chat/completions", + ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 55e635f3e17..4a61fd30fc9 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -96,6 +96,19 @@ def test_supports_function_calling_github_anthropic_alias(): ) +def test_supports_function_calling_deepinfra_llama(): + """Test that deepinfra Llama models correctly report function calling support. + + Regression test for https://github.com/BerriAI/litellm/issues/22619 + """ + assert ( + litellm.utils.supports_function_calling( + model="deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo" + ) + is True + ) + + def test_supports_function_calling_unknown_github_alias_returns_false(): assert ( litellm.utils.supports_function_calling( @@ -2914,6 +2927,38 @@ class TestIsCachedMessage: message = {"role": "user", "content": []} assert is_cached_message(message) is False + def test_message_level_cache_control_returns_true(self): + """Message with string content and message-level cache_control should return True. + + This is the format injected by the cache_control_injection_points hook + when the message content is a string (common for system messages). + Fixes GitHub issue #18519 - Gemini models ignoring cache_control_injection_points. + """ + message = { + "role": "system", + "content": "You are a helpful assistant.", + "cache_control": {"type": "ephemeral"}, + } + assert is_cached_message(message) is True + + def test_message_level_cache_control_wrong_type_returns_false(self): + """Message-level cache_control with non-ephemeral type should return False.""" + message = { + "role": "system", + "content": "You are a helpful assistant.", + "cache_control": {"type": "permanent"}, + } + assert is_cached_message(message) is False + + def test_message_level_cache_control_non_dict_returns_false(self): + """Message-level cache_control that's not a dict should return False.""" + message = { + "role": "system", + "content": "You are a helpful assistant.", + "cache_control": "ephemeral", + } + assert is_cached_message(message) is False + @pytest.mark.asyncio class TestProxyLoggingBudgetAlerts: diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx index a74d3c108d6..b3829d0a8f4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx @@ -374,6 +374,27 @@ const Sidebar2: React.FC = ({ accessToken, userRole, defaultSelect router.push(href); }; + // Wrap label in so every nav item supports right-click → "Open in new tab" + // and Ctrl/Cmd+click to open in a new tab, while preserving SPA navigation for normal clicks. + const renderNavLink = (label: string, page: string): React.ReactNode => { + const href = toHref(page); + return ( + { + if (e.metaKey || e.ctrlKey || e.shiftKey || e.button === 1) { + e.stopPropagation(); + return; + } + e.preventDefault(); + }} + style={{ color: "inherit", textDecoration: "none" }} + > + {label} + + ); + }; + return ( = ({ accessToken, userRole, defaultSelect items={filteredMenuItems.map((item) => ({ key: item.key, icon: item.icon, - label: item.label, + label: renderNavLink(item.label, item.page), children: item.children?.map((child) => ({ key: child.key, icon: child.icon, - label: child.label, + label: renderNavLink(child.label, child.page), onClick: () => goTo(child.page), })), onClick: !item.children ? () => goTo(item.page) : undefined, diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 622c3bf70a9..b927f312df8 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -39,7 +39,7 @@ import UserDashboard from "@/components/user_dashboard"; import { AccessGroupsPage } from "@/components/AccessGroups/AccessGroupsPage"; import { ProjectsPage } from "@/components/Projects/ProjectsPage"; import VectorStoreManagement from "@/components/vector_store_management"; -import ToolPolicies from "@/components/ToolPolicies"; +import ToolPoliciesView from "@/components/ToolPoliciesView"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; @@ -549,7 +549,7 @@ function CreateKeyPageContent() { ) : page == "vector-stores" ? ( ) : page == "tool-policies" ? ( - + ) : page == "guardrails-monitor" ? ( ) : page == "new_usage" ? ( diff --git a/ui/litellm-dashboard/src/components/ToolDetail.tsx b/ui/litellm-dashboard/src/components/ToolDetail.tsx new file mode 100644 index 00000000000..ed0f866acb8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ToolDetail.tsx @@ -0,0 +1,445 @@ +"use client"; + +import { ArrowLeftOutlined, HistoryOutlined, ToolOutlined } from "@ant-design/icons"; +import { useQuery, useQueryClient } from "@tanstack/react-query"; +import { Button, Select, Spin } from "antd"; +import React, { useCallback, useEffect, useMemo, useState } from "react"; +import TeamDropdown from "@/components/common_components/team_dropdown"; +import { LogViewer } from "@/components/GuardrailsMonitor/LogViewer"; +import type { LogEntry } from "@/components/GuardrailsMonitor/mockData"; +import { PolicySelect } from "@/components/ToolPolicies/PolicySelect"; +import { + deleteToolPolicyOverride, + fetchToolDetail, + fetchToolPolicyOptions, + getToolUsageLogs, + keyListCall, + teamListCall, + updateToolPolicy, + type ToolPolicyOption, + type ToolPolicyOverrideRow, +} from "@/components/networking"; +import type { Team } from "@/components/key_team_helpers/key_list"; + +interface ToolDetailProps { + toolName: string; + onBack: () => void; + accessToken: string | null; +} + +interface KeyOption { + token: string; + key_alias?: string; +} + +const TOOL_DETAIL_QUERY_KEY = "tool-detail"; + +const LOGS_PAGE_SIZE = 50; + +function getDefaultLogsDateRange(): { start: string; end: string } { + const end = new Date(); + const start = new Date(); + start.setDate(start.getDate() - 90); + const fmt = (d: Date) => + d.toISOString().slice(0, 19).replace("T", " "); + return { start: fmt(start), end: fmt(end) }; +} + +export function ToolDetail({ toolName, onBack, accessToken }: ToolDetailProps) { + const queryClient = useQueryClient(); + const [overrideSaving, setOverrideSaving] = useState(false); + const [inputPolicySaving, setInputPolicySaving] = useState(false); + const [outputPolicySaving, setOutputPolicySaving] = useState(false); + const [blockScope, setBlockScope] = useState<"team" | "key">("team"); + const [blockTeamId, setBlockTeamId] = useState(null); + const [blockKey, setBlockKey] = useState(null); + + const logsDateRange = useMemo(() => getDefaultLogsDateRange(), []); + + const { data: detail, isLoading: detailLoading, error: detailError } = useQuery({ + queryKey: [TOOL_DETAIL_QUERY_KEY, toolName], + queryFn: () => fetchToolDetail(accessToken!, toolName), + enabled: !!accessToken && !!toolName, + }); + + const { data: policyOptions } = useQuery({ + queryKey: ["tool-policy-options"], + queryFn: () => fetchToolPolicyOptions(accessToken!), + enabled: !!accessToken, + staleTime: 60_000, + }); + + const { data: teamsData } = useQuery({ + queryKey: ["teams-list-tool-detail"], + queryFn: () => teamListCall(accessToken!, null, null), + enabled: !!accessToken, + }); + + const { data: keysData } = useQuery({ + queryKey: ["keys-list-tool-detail"], + queryFn: () => keyListCall(accessToken!, null, null, null, null, null, 1, 100), + enabled: !!accessToken, + }); + + const { data: logsData, isLoading: logsLoading } = useQuery({ + queryKey: ["tool-usage-logs", toolName, logsDateRange.start, logsDateRange.end], + queryFn: () => + getToolUsageLogs(accessToken!, toolName, { + page: 1, + pageSize: LOGS_PAGE_SIZE, + startDate: logsDateRange.start, + endDate: logsDateRange.end, + }), + enabled: !!accessToken && !!toolName, + }); + + const logs: LogEntry[] = useMemo(() => { + const list = logsData?.logs ?? []; + return list.map((l) => ({ + id: l.id, + timestamp: l.timestamp, + action: "passed" as const, + model: l.model ?? undefined, + input_snippet: l.input_snippet ?? undefined, + })); + }, [logsData?.logs]); + + const teams: Team[] = useMemo(() => { + const arr = Array.isArray(teamsData) ? teamsData : teamsData?.data ?? []; + return arr.map((t: { team_id?: string; id?: string; team_alias?: string }) => ({ + team_id: t.team_id ?? t.id ?? "", + team_alias: t.team_alias ?? t.team_id ?? "", + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: "", + created_at: "", + keys: [], + members_with_roles: [], + spend: 0, + })); + }, [teamsData]); + + const keys: KeyOption[] = useMemo(() => { + const keysRes = keysData?.keys ?? keysData?.data ?? []; + return keysRes.map((k: { token?: string; api_key?: string; key_hash?: string; key_alias?: string }) => ({ + token: k.token ?? k.api_key ?? k.key_hash ?? "", + key_alias: k.key_alias ?? (k.token ?? k.api_key ?? k.key_hash)?.toString?.()?.substring?.(0, 8), + })); + }, [keysData]); + + const invalidateDetail = useCallback(() => { + queryClient.invalidateQueries({ queryKey: [TOOL_DETAIL_QUERY_KEY, toolName] }); + }, [queryClient, toolName]); + + const handleInputPolicyChange = useCallback( + async (_name: string, newPolicy: string) => { + if (!accessToken) return; + setInputPolicySaving(true); + try { + await updateToolPolicy(accessToken, toolName, { input_policy: newPolicy }); + invalidateDetail(); + } catch (e: unknown) { + alert(`Failed to update input policy: ${e instanceof Error ? e.message : String(e)}`); + } finally { + setInputPolicySaving(false); + } + }, + [accessToken, toolName, invalidateDetail] + ); + + const handleOutputPolicyChange = useCallback( + async (_name: string, newPolicy: string) => { + if (!accessToken) return; + setOutputPolicySaving(true); + try { + await updateToolPolicy(accessToken, toolName, { output_policy: newPolicy }); + invalidateDetail(); + } catch (e: unknown) { + alert(`Failed to update output policy: ${e instanceof Error ? e.message : String(e)}`); + } finally { + setOutputPolicySaving(false); + } + }, + [accessToken, toolName, invalidateDetail] + ); + + const handleAddOverride = useCallback(async () => { + if (!accessToken || !toolName) return; + const isTeam = blockScope === "team"; + if (isTeam && !blockTeamId) return; + if (!isTeam && !blockKey?.token) return; + setOverrideSaving(true); + try { + await updateToolPolicy(accessToken, toolName, { input_policy: "blocked" }, { + team_id: isTeam ? blockTeamId : undefined, + key_hash: !isTeam ? blockKey!.token : undefined, + key_alias: !isTeam ? blockKey!.key_alias : undefined, + }); + invalidateDetail(); + setBlockTeamId(null); + setBlockKey(null); + } catch (e: unknown) { + alert(`Failed to add override: ${e instanceof Error ? e.message : String(e)}`); + } finally { + setOverrideSaving(false); + } + }, [accessToken, toolName, blockScope, blockTeamId, blockKey, invalidateDetail]); + + const handleRemoveOverride = useCallback( + async (override: ToolPolicyOverrideRow) => { + if (!accessToken || !toolName) return; + setOverrideSaving(true); + try { + await deleteToolPolicyOverride(accessToken, toolName, { + team_id: override.team_id ?? undefined, + key_hash: override.key_hash ?? undefined, + }); + invalidateDetail(); + } catch (e: unknown) { + alert(`Failed to remove override: ${e instanceof Error ? e.message : String(e)}`); + } finally { + setOverrideSaving(false); + } + }, + [accessToken, toolName, invalidateDetail] + ); + + if (detailLoading && !detail) { + return ( +
+ +
+ ); + } + + if (detailError && !detail) { + return ( +
+ +

Failed to load tool details.

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

{tool.tool_name}

+ + {tool.origin ?? "—"} + + + {(tool.call_count ?? 0).toLocaleString()} calls + +
+
+ {tool.user_agent && ( +
+
User Agent:
+
{tool.user_agent}
+
+ )} + {tool.created_at && ( +
+
First Discovered:
+
{new Date(tool.created_at).toLocaleString()}
+
+ )} + {tool.last_used_at && ( +
+
Last Used:
+
{new Date(tool.last_used_at).toLocaleString()}
+
+ )} +
+
+
+
+ +
+ {/* Two-panel policy layout */} +
+
+

Input Policy

+

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

+ +
+ +
+

Output Policy

+

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

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

Blocked for team or key

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

Block for team or key

+
+
+ Scope +
+ + +
+
+
+ + {blockScope === "team" ? "Team" : "Key"} + + {blockScope === "team" ? ( + setBlockTeamId(id || null)} + /> + ) : ( + onChange(toolName, v)} - onClick={(e) => e.stopPropagation()} - style={{ - minWidth: 110, - fontWeight: 500, - }} - popupMatchSelectWidth={false} - options={POLICY_OPTIONS.map((o) => ({ - value: o.value, - label: ( - - - {o.label} - - ), - }))} - /> - ); -}; - -export const ToolPolicies: React.FC = ({ accessToken }) => { +export const ToolPolicies: React.FC = ({ accessToken, onSelectTool }) => { const [tools, setTools] = useState([]); const [loading, setLoading] = useState(true); const [isFetching, setIsFetching] = useState(false); const [error, setError] = useState(null); - const [saving, setSaving] = useState(null); + const [savingInput, setSavingInput] = useState(null); + const [savingOutput, setSavingOutput] = useState(null); const [searchTerm, setSearchTerm] = useState(""); const [sortField, setSortField] = useState("created_at"); @@ -123,16 +96,29 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { return () => clearInterval(id); }, [isLiveTail, load]); - const handlePolicyChange = async (toolName: string, newPolicy: string) => { + const handleInputPolicyChange = async (toolName: string, newPolicy: string) => { if (!accessToken) return; - setSaving(toolName); + setSavingInput(toolName); try { - await updateToolPolicy(accessToken, toolName, newPolicy); - setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t))); + await updateToolPolicy(accessToken, toolName, { input_policy: newPolicy }); + setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, input_policy: newPolicy } : t))); } catch (e: any) { - alert(`Failed to update policy: ${e.message}`); + alert(`Failed to update input policy: ${e.message}`); } finally { - setSaving(null); + setSavingInput(null); + } + }; + + const handleOutputPolicyChange = async (toolName: string, newPolicy: string) => { + if (!accessToken) return; + setSavingOutput(toolName); + try { + await updateToolPolicy(accessToken, toolName, { output_policy: newPolicy }); + setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, output_policy: newPolicy } : t))); + } catch (e: any) { + alert(`Failed to update output policy: ${e.message}`); + } finally { + setSavingOutput(null); } }; @@ -157,7 +143,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { setCurrentPage(1); }; - // Build unique team/key options from loaded data const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map((v) => ({ label: v as string, value: v as string, @@ -169,9 +154,14 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { const filterOptions: FilterOption[] = [ { - name: "Policy", - label: "Policy", - options: POLICY_OPTIONS.map((o) => ({ label: o.label, value: o.value })), + name: "Input Policy", + label: "Input Policy", + options: INPUT_POLICY_OPTIONS.map((o) => ({ label: o.label, value: o.value })), + }, + { + name: "Output Policy", + label: "Output Policy", + options: OUTPUT_POLICY_OPTIONS.map((o) => ({ label: o.label, value: o.value })), }, { name: "Team Name", @@ -185,6 +175,39 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { }, ]; + const { newToday, newYesterday, trendSubtitle, totalTools, blockedCount, activeTeamsCount, needsReviewTools } = + useMemo(() => { + const now = new Date(); + const todayKey = getUTCDateKey(now); + const yesterday = new Date(now); + yesterday.setUTCDate(yesterday.getUTCDate() - 1); + const yesterdayKey = getUTCDateKey(yesterday); + + const newToday = countToolsInUTCDay(tools, todayKey); + const newYesterday = countToolsInUTCDay(tools, yesterdayKey); + const trendSubtitle = getTrendSubtitle(newToday, newYesterday); + + const totalTools = tools.length; + const blockedCount = tools.filter((t) => t.input_policy === "blocked").length; + const activeTeamsCount = new Set(tools.map((t) => t.team_id).filter(Boolean)).size; + + const needsReviewTools = tools.filter( + (t) => + isCreatedInUTCDay(t.created_at, todayKey) && + t.input_policy === "untrusted" + ); + + return { + newToday, + newYesterday, + trendSubtitle, + totalTools, + blockedCount, + activeTeamsCount, + needsReviewTools, + }; + }, [tools]); + const SortHeader = ({ label, field }: { label: string; field: SortField }) => (
{label} @@ -203,10 +226,12 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { (t.team_id ?? "").toLowerCase().includes(q) || (t.key_alias ?? "").toLowerCase().includes(q) || (t.key_hash ?? "").toLowerCase().includes(q) || - t.call_policy.toLowerCase().includes(q); + t.input_policy.toLowerCase().includes(q) || + t.output_policy.toLowerCase().includes(q); if (!matchesSearch) return false; } - if (activeFilters["Policy"] && t.call_policy !== activeFilters["Policy"]) return false; + if (activeFilters["Input Policy"] && t.input_policy !== activeFilters["Input Policy"]) return false; + if (activeFilters["Output Policy"] && t.output_policy !== activeFilters["Output Policy"]) return false; if (activeFilters["Team Name"] && t.team_id !== activeFilters["Team Name"]) return false; if (activeFilters["Key Name"] && t.key_alias !== activeFilters["Key Name"]) return false; return true; @@ -223,11 +248,74 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { const totalPages = Math.max(1, Math.ceil(sorted.length / pageSize)); const paginated = sorted.slice((currentPage - 1) * pageSize, currentPage * pageSize); + const scrollToToolRow = (toolId: string) => { + const idx = sorted.findIndex((t) => t.tool_id === toolId); + if (idx >= 0) { + const page = Math.floor(idx / pageSize) + 1; + if (page !== currentPage) setCurrentPage(page); + requestAnimationFrame(() => { + setTimeout(() => { + document.getElementById(`tool-row-${toolId}`)?.scrollIntoView({ behavior: "smooth", block: "center" }); + }, 100); + }); + } + }; + return ( -
+

Tool Policies

+ +
+ + + + } + /> + + 0 ? "text-red-600" : undefined} + /> + 0 ? activeTeamsCount : "—"} /> +
+ + {needsReviewTools.length > 0 && ( +
+

Needs Review

+

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

+
+ {needsReviewTools.map((t) => ( + + + {t.tool_name} + + + + ))} +
+
+ )} +
- {/* Toolbar */}
@@ -311,7 +399,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
- {/* Filter row */}
= ({ accessToken }) => {
- {/* Auto-refresh banner */} {isLiveTail && (
Auto-refreshing every 15 seconds @@ -336,7 +422,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
{error}
)} - {/* Table */} @@ -347,7 +432,10 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { - + + + + @@ -359,45 +447,61 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { - Origin + User Agent {loading ? ( - + Loading tools… ) : paginated.length === 0 ? ( - + No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery. ) : ( paginated.map((tool) => ( - + - - - {tool.tool_name} - - + - - {(tool.call_count ?? 0).toLocaleString()} + + + + +
+ {(tool.call_count ?? 0).toLocaleString()} +
@@ -417,8 +521,8 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { - - {tool.origin ?? "-"} + + {tool.user_agent ?? "-"} @@ -427,7 +531,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
- {/* Bottom pagination (only when > 1 page) */} {totalPages > 1 && (
@@ -453,6 +556,7 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
)}
+
); }; diff --git a/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.tsx b/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.tsx new file mode 100644 index 00000000000..1317351931e --- /dev/null +++ b/ui/litellm-dashboard/src/components/ToolPolicies/PolicySelect.tsx @@ -0,0 +1,92 @@ +"use client"; + +import React from "react"; +import { Select } from "antd"; + +export const INPUT_POLICY_OPTIONS = [ + { value: "untrusted", label: "untrusted", color: "#92400e", bg: "#fef3c7", border: "#fcd34d" }, + { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, + { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, +] as const; + +export const OUTPUT_POLICY_OPTIONS = [ + { value: "untrusted", label: "untrusted", color: "#92400e", bg: "#fef3c7", border: "#fcd34d" }, + { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, +] as const; + +export const POLICY_OPTIONS = INPUT_POLICY_OPTIONS; + +export const policyStyle = (p: string) => + INPUT_POLICY_OPTIONS.find((o) => o.value === p) ?? INPUT_POLICY_OPTIONS[0]; + +export interface PolicySelectProps { + value: string; + toolName: string; + saving: boolean; + onChange: (toolName: string, policy: string) => void; + policyType?: "input" | "output"; + size?: "small" | "middle"; + minWidth?: number; + stopPropagation?: boolean; +} + +export const PolicySelect: React.FC = ({ + value, + toolName, + saving, + onChange, + policyType = "input", + size = "small", + minWidth = 110, + stopPropagation = true, +}) => { + const options = policyType === "output" ? OUTPUT_POLICY_OPTIONS : INPUT_POLICY_OPTIONS; + const style = policyStyle(value); + return ( +