From 48fe111c12f6599fafb694946d86c9d1e1dd77a7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:47:48 +0000 Subject: [PATCH 001/102] fix(responses): stop managed Responses WS from leaking litellm_params into provider body Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 3 +- .../test_responses_websocket_all_providers.py | 53 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index eb78e6f9c8d..616a6659a55 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2097,8 +2097,7 @@ class ManagedResponsesWebSocketHandler: if "litellm_metadata" not in call_kwargs: call_kwargs["litellm_metadata"] = {} call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request - call_kwargs.setdefault("litellm_params", {}) - call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request + call_kwargs["proxy_server_request"] = proxy_server_request async def _stream_and_forward(self, model: str, call_kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]: """ diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 4509abc7749..7557757ded1 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -660,6 +660,59 @@ class TestChunkTransformation: assert ManagedResponsesWebSocketHandler._input_to_messages({}) == [] +class TestUpdateProxyRequest: + """Regression tests for ManagedResponsesWebSocketHandler._update_proxy_request. + + The managed WebSocket path calls ``litellm.aresponses(model=..., **call_kwargs)``. + ``litellm_params`` is not a Responses API request field, so passing it as a + top-level kwarg leaks it into the provider request body and providers that + forbid extra inputs (e.g. Anthropic) reject the call with + ``litellm_params: Extra inputs are not permitted``. The request-tracking data + must ride along as ``proxy_server_request`` instead, which litellm consumes + internally and never forwards to the provider. + """ + + def test_does_not_inject_litellm_params_kwarg(self): + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + call_kwargs = { + "input": "hello", + "store": True, + "litellm_metadata": { + "proxy_server_request": {"headers": {}, "body": {}}, + }, + } + + ManagedResponsesWebSocketHandler._update_proxy_request( + call_kwargs, "anthropic/claude-sonnet-4-5" + ) + + assert "litellm_params" not in call_kwargs + assert call_kwargs["proxy_server_request"]["body"]["model"] == ( + "anthropic/claude-sonnet-4-5" + ) + assert call_kwargs["proxy_server_request"]["body"]["input"] == "hello" + + def test_proxy_server_request_matches_metadata(self): + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + call_kwargs = { + "input": "hi", + "litellm_metadata": {"proxy_server_request": {"body": {}}}, + } + + ManagedResponsesWebSocketHandler._update_proxy_request(call_kwargs, "gpt-4o") + + assert ( + call_kwargs["proxy_server_request"] + == call_kwargs["litellm_metadata"]["proxy_server_request"] + ) + + class TestWebSocketEventTypes: """Test that all WebSocket event types are properly handled with dict-based chunks""" From 9c541d9ce3ecf31ca8673ebc4472250f477b55ae Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:27:33 +0000 Subject: [PATCH 002/102] fix(ui): show internal user email in logs table and log detail drawer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../(dashboard)/hooks/users/useUsers.test.ts | 61 ++++++++++++++++++- .../app/(dashboard)/hooks/users/useUsers.ts | 18 ++++++ .../LogDetailContent.integration.test.tsx | 23 +++++++ .../LogDetailsDrawer/LogDetailContent.tsx | 23 ++++++- .../LogDetailsDrawer.test.tsx | 39 +++++++++++- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 3 + .../components/view_logs/RequestLogsTable.tsx | 10 ++- .../RequestLogsTableColumns.test.tsx | 21 +++++++ .../view_logs/RequestLogsTableColumns.tsx | 23 ++++++- 9 files changed, 212 insertions(+), 9 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts index dd7209140a9..2e8471ba84f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts @@ -2,7 +2,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { renderHook, waitFor } from "@testing-library/react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import React, { ReactNode } from "react"; -import { useInfiniteUsers, useUserLookup } from "./useUsers"; +import { useInfiniteUsers, useUserEmailLookup, useUserLookup } from "./useUsers"; import { userListCall } from "@/components/networking"; import type { UserListResponse } from "@/components/networking"; @@ -335,3 +335,62 @@ describe("useUserLookup", () => { expect(userListCall).not.toHaveBeenCalled(); }); }); + +describe("useUserEmailLookup", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue(DEFAULT_AUTH); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("fetches the distinct ids in one call and maps each id to its email", async () => { + const response = buildUserListResponse(1, 1, 2); + vi.mocked(userListCall).mockResolvedValue(response); + + const { result } = renderHook(() => useUserEmailLookup(["user-1-1", "user-1-0", "user-1-1", ""]), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(userListCall).toHaveBeenCalledTimes(1); + expect(userListCall).toHaveBeenCalledWith("test-access-token", ["user-1-0", "user-1-1"], 1, 2); + expect(result.current.data).toEqual({ + "user-1-0": "user-1-0@example.com", + "user-1-1": "user-1-1@example.com", + }); + }); + + it("omits users that have no email so callers fall back to the id", async () => { + const response = buildUserListResponse(1, 1, 2); + vi.mocked(userListCall).mockResolvedValue({ + ...response, + users: [{ ...response.users[0], user_email: "" }, response.users[1]], + }); + + const { result } = renderHook(() => useUserEmailLookup(["user-1-0", "user-1-1"]), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(result.current.data).toEqual({ "user-1-1": "user-1-1@example.com" }); + }); + + it("does not query with no ids", async () => { + const { result } = renderHook(() => useUserEmailLookup([]), { wrapper }); + + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(result.current.fetchStatus).toBe("idle"); + expect(userListCall).not.toHaveBeenCalled(); + }); + + it("does not query for a non-admin role", async () => { + mockUseAuthorized.mockReturnValue({ ...DEFAULT_AUTH, userRole: "Internal User" }); + + const { result } = renderHook(() => useUserEmailLookup(["user-1-0"]), { wrapper }); + + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(result.current.fetchStatus).toBe("idle"); + expect(userListCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts index 011e43777b5..4a28e7ff3f2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts @@ -38,6 +38,24 @@ export const useInfiniteUsers = (pageSize: number = DEFAULT_PAGE_SIZE, searchEma }); }; +const USER_LIST_MAX_PAGE_SIZE = 100; + +export const useUserEmailLookup = (userIds: readonly string[]) => { + const { accessToken, userRole } = useAuthorized(); + const distinctIds = Array.from(new Set(userIds.filter((id) => id !== ""))).sort(); + return useQuery>({ + queryKey: userLookupKeys.list({ filters: { ids: distinctIds.join(",") } }), + queryFn: async () => { + const ids = distinctIds.slice(0, USER_LIST_MAX_PAGE_SIZE); + const response = await userListCall(accessToken!, ids, 1, ids.length); + return Object.fromEntries( + response.users.filter((user) => Boolean(user.user_email)).map((user) => [user.user_id, user.user_email]), + ); + }, + enabled: Boolean(accessToken) && distinctIds.length > 0 && all_admin_roles.includes(userRole!), + }); +}; + export const useUserLookup = (userId: string | null) => { const { accessToken, userRole } = useAuthorized(); return useQuery({ diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx index 721525268b3..793bfa7c246 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx @@ -56,6 +56,29 @@ describe("LogDetailContent", () => { expect(screen.getByText("completion")).toBeInTheDocument(); }); + it("shows the requesting user's email and id in Request Details when the email is resolved", () => { + render( + , + ); + + expect(screen.getByText("User")).toBeInTheDocument(); + expect(screen.getByText("alice@example.com")).toBeInTheDocument(); + expect(screen.getByText("106514937785257944828")).toBeInTheDocument(); + }); + + it("falls back to the user id in Request Details when no email is resolved", () => { + render(); + + expect(screen.getByText("User")).toBeInTheDocument(); + expect(screen.getByText("106514937785257944828")).toBeInTheDocument(); + }); + + it("omits the User row when the log has no internal user", () => { + render(); + + expect(screen.queryByText("User")).not.toBeInTheDocument(); + }); + it("should display error alert when request has failed", () => { render( {logEntry.model} {logEntry.custom_llm_provider || "-"} {logEntry.call_type} + {logEntry.user && ( + + + + )} @@ -333,6 +344,16 @@ function TagsSection({ tags }: { tags: Record }) { ); } +function UserIdentity({ userId, email }: { userId: string; email?: string }) { + if (!email || email === userId) return ; + return ( + + {email} + + + ); +} + function GuardrailLabel({ label, maskedCount }: { label: string; maskedCount: number }) { const handleClick = () => { const el = document.getElementById("guardrail-section"); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.test.tsx index b96e5279722..f5bf7cda951 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.test.tsx @@ -14,8 +14,13 @@ vi.mock("@/app/(dashboard)/hooks/logDetails/useLogDetails", () => ({ useLogDetails: () => ({ data: null, isLoading: false }), })); +const mockUseUserLookup = vi.fn(() => ({ data: undefined })); +vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({ + useUserLookup: (userId: string | null) => mockUseUserLookup(userId), +})); + vi.mock("./LogDetailContent", () => ({ - LogDetailContent: () => null, + LogDetailContent: ({ userEmail }: { userEmail?: string }) => user-email:{userEmail ?? "none"}, GuardrailJumpLink: () => null, })); @@ -124,6 +129,38 @@ describe("LogDetailsDrawer session sidebar sorting", () => { }); }); +describe("LogDetailsDrawer internal user email", () => { + const renderSingleLog = (user: string | undefined) => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + {}} + logEntry={makeLog({ request_id: "single", user })} + accessToken="token" + /> + , + ); + }; + + it("looks up the log's internal user and hands the resolved email to the detail content", () => { + mockUseUserLookup.mockReturnValue({ data: { user_id: "u-1", user_email: "alice@example.com" } }); + renderSingleLog("u-1"); + + expect(mockUseUserLookup).toHaveBeenCalledWith("u-1"); + expect(screen.getByText("user-email:alice@example.com")).toBeInTheDocument(); + }); + + it("skips the lookup and passes no email when the log has no internal user", () => { + mockUseUserLookup.mockReturnValue({ data: undefined }); + renderSingleLog(undefined); + + expect(mockUseUserLookup).toHaveBeenCalledWith(null); + expect(screen.getByText("user-email:none")).toBeInTheDocument(); + }); +}); + describe("LogDetailsDrawer session sidebar auto-router icon", () => { const routedSessionLogs = [ makeLog({ request_id: "routed", model: "claude-opus-4-8", model_group: "smart-router" }), diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index ddd0a650c04..dc9207d59eb 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -17,6 +17,7 @@ import { getSpendString } from "@/utils/dataUtils"; import { normalizeGuardrailEntries, sortSessionLogs, SessionLogSortMode } from "./utils"; import { DRAWER_WIDTH } from "./constants"; import { useLogDetails } from "@/app/(dashboard)/hooks/logDetails/useLogDetails"; +import { useUserLookup } from "@/app/(dashboard)/hooks/users/useUsers"; export interface LogDetailsDrawerProps { open: boolean; @@ -245,6 +246,7 @@ export function LogDetailsDrawer({ const logDetails = useLogDetails(currentLog?.request_id, startTime, open && !!currentLog?.request_id); const detailsData = logDetails.data as any; const isLoadingDetails = logDetails.isLoading; + const { data: logUser } = useUserLookup(open && currentLog?.user ? currentLog.user : null); // Build an enriched log entry that merges lazy-loaded details. // The list endpoint may already include messages/response when store_prompts_in_spend_logs is enabled, @@ -465,6 +467,7 @@ export function LogDetailsDrawer({ logEntry={enrichedLog} isLoadingDetails={isLoadingDetails} accessToken={accessToken ?? null} + userEmail={logUser?.user_email || undefined} /> diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx index c3d204e1ae3..5a9bac428c8 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx @@ -4,6 +4,7 @@ import type { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } fr import { ScrollText } from "lucide-react"; import { useMemo, useState, type ReactNode } from "react"; +import { useUserEmailLookup } from "@/app/(dashboard)/hooks/users/useUsers"; import { DataTable, DataTableFilterDrawer, DataTableToolbar } from "@/components/shared/DataTable"; import type { Team } from "../key_team_helpers/key_list"; @@ -73,10 +74,13 @@ export function RequestLogsTable({ }: RequestLogsTableProps) { const [filtersOpen, setFiltersOpen] = useState(false); + const userIds = useMemo(() => data.flatMap((log) => (log.user ? [log.user] : [])), [data]); + const { data: emailByUserId } = useUserEmailLookup(userIds); + const columns = useMemo(() => { - const deps = { onKeyHashClick, onSessionClick }; - return getRequestLogsTableColumns(deps); - }, [onKeyHashClick, onSessionClick]); + const resolveUserEmail = (userId: string) => emailByUserId?.[userId]; + return getRequestLogsTableColumns({ onKeyHashClick, onSessionClick, resolveUserEmail }); + }, [onKeyHashClick, onSessionClick, emailByUserId]); const isFiltered = columnFilters.length > 0 || searchValue !== ""; diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index 9f0e659cb1f..08853814269 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -75,6 +75,27 @@ describe("Cost column", () => { }); }); +describe("Internal User column", () => { + const emailById: Record = { "106514937785257944828": "alice@example.com" }; + const deps = { ...noopDeps, resolveUserEmail: (userId: string) => emailById[userId] }; + + it("shows the user's email instead of the raw id, with both in the tooltip", async () => { + const user = userEvent.setup(); + renderRows([logEntry({ request_id: "req-known-user", user: "106514937785257944828" })], deps); + + const emailCell = screen.getByText("alice@example.com"); + expect(screen.queryByText("106514937785257944828")).not.toBeInTheDocument(); + await user.hover(emailCell); + expect(await screen.findByText("alice@example.com (106514937785257944828)")).toBeInTheDocument(); + }); + + it("falls back to the raw id when no email is known for the user", () => { + renderRows([logEntry({ request_id: "req-unknown-user", user: "unknown-user-id" })], deps); + + expect(screen.getByText("unknown-user-id")).toBeInTheDocument(); + }); +}); + describe("Tokens column", () => { const sessionRow: Partial = { request_id: "req-session-tokens", diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index 1ec1087a1a4..dd83ad6eb05 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -15,6 +15,7 @@ import { AgentBadge, AgentIcon, BatchBadge, LlmBadge, McpBadge, SparkleIcon, Wre export interface RequestLogsTableColumnsDeps { onKeyHashClick: (keyHash: string) => void; onSessionClick: (log: LogEntry) => void; + resolveUserEmail?: (userId: string) => string | undefined; } const readMetaString = (metadata: Record | undefined, key: string): string | undefined => { @@ -32,14 +33,25 @@ const readMcpLogoUrl = (metadata: Record | undefined): string | const getLogoUrl = (row: LogEntry, provider: string): string => readMcpLogoUrl(row.metadata) ?? (provider ? getProviderLogoAndName(provider).logo : ""); -function TruncatedText({ value }: { value: string | undefined }) { +function TruncatedText({ value, tooltip }: { value: string | undefined; tooltip?: string }) { const display = value ?? "-"; - return {display}} />; + return ( + {display}} + /> + ); +} + +function UserCell({ userId, email }: { userId: string | undefined; email: string | undefined }) { + if (!userId || !email || email === userId) return ; + return ; } export const getRequestLogsTableColumns = ({ onKeyHashClick, onSessionClick, + resolveUserEmail = () => undefined, }: RequestLogsTableColumnsDeps): ColumnDef[] => [ { id: "startTime", @@ -313,7 +325,12 @@ export const getRequestLogsTableColumns = ({ header: "Internal User", size: 150, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => ( + + ), }, { id: "end_user", From bde96e3197bf1d1d9af07fa238f63649b543adf0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:36:29 +0000 Subject: [PATCH 003/102] fix(ui): keep user id boundaries in email lookup query key Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/app/(dashboard)/hooks/users/useUsers.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts index 4a28e7ff3f2..84aeba90ae2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts @@ -44,7 +44,7 @@ export const useUserEmailLookup = (userIds: readonly string[]) => { const { accessToken, userRole } = useAuthorized(); const distinctIds = Array.from(new Set(userIds.filter((id) => id !== ""))).sort(); return useQuery>({ - queryKey: userLookupKeys.list({ filters: { ids: distinctIds.join(",") } }), + queryKey: userLookupKeys.list({ filters: { ids: JSON.stringify(distinctIds) } }), queryFn: async () => { const ids = distinctIds.slice(0, USER_LIST_MAX_PAGE_SIZE); const response = await userListCall(accessToken!, ids, 1, ids.length); From 3630642110e0055040f5c5666cb2ff1c195d519e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:21:20 +0000 Subject: [PATCH 004/102] fix(ui): allow Org Admin session role to resolve user emails in logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/app/(dashboard)/hooks/users/useUsers.test.ts | 12 +++++++++++- .../src/app/(dashboard)/hooks/users/useUsers.ts | 8 ++++---- ui/litellm-dashboard/src/utils/roles.ts | 6 ++++++ 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts index 2e8471ba84f..f49c728446b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.test.ts @@ -235,7 +235,7 @@ describe("useInfiniteUsers", () => { }); it("should execute query for each admin role", async () => { - const adminRoles = ["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer", "org_admin"]; + const adminRoles = ["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer", "org_admin", "Org Admin"]; for (const role of adminRoles) { vi.clearAllMocks(); @@ -384,6 +384,16 @@ describe("useUserEmailLookup", () => { expect(userListCall).not.toHaveBeenCalled(); }); + it("queries for the formatted Org Admin session role", async () => { + mockUseAuthorized.mockReturnValue({ ...DEFAULT_AUTH, userRole: "Org Admin" }); + vi.mocked(userListCall).mockResolvedValue(buildUserListResponse(1, 1, 1)); + + const { result } = renderHook(() => useUserEmailLookup(["user-1-0"]), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(result.current.data).toEqual({ "user-1-0": "user-1-0@example.com" }); + }); + it("does not query for a non-admin role", async () => { mockUseAuthorized.mockReturnValue({ ...DEFAULT_AUTH, userRole: "Internal User" }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts index 84aeba90ae2..3b7f9fbeb02 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useUsers.ts @@ -1,7 +1,7 @@ import { userListCall, UserInfo, UserListResponse } from "@/components/networking"; import { useInfiniteQuery, useQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; -import { all_admin_roles } from "@/utils/roles"; +import { canListUsers } from "@/utils/roles"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; const infiniteUsersKeys = createQueryKeys("infiniteUsers"); @@ -34,7 +34,7 @@ export const useInfiniteUsers = (pageSize: number = DEFAULT_PAGE_SIZE, searchEma } return undefined; }, - enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!), + enabled: Boolean(accessToken) && canListUsers(userRole), }); }; @@ -52,7 +52,7 @@ export const useUserEmailLookup = (userIds: readonly string[]) => { response.users.filter((user) => Boolean(user.user_email)).map((user) => [user.user_id, user.user_email]), ); }, - enabled: Boolean(accessToken) && distinctIds.length > 0 && all_admin_roles.includes(userRole!), + enabled: Boolean(accessToken) && distinctIds.length > 0 && canListUsers(userRole), }); }; @@ -64,6 +64,6 @@ export const useUserLookup = (userId: string | null) => { const response = await userListCall(accessToken!, [userId!], 1, 1); return response.users.find((user) => user.user_id === userId) ?? null; }, - enabled: Boolean(accessToken) && Boolean(userId) && all_admin_roles.includes(userRole!), + enabled: Boolean(accessToken) && Boolean(userId) && canListUsers(userRole), }); }; diff --git a/ui/litellm-dashboard/src/utils/roles.ts b/ui/litellm-dashboard/src/utils/roles.ts index 62a5f02cc39..1cb7e75c19b 100644 --- a/ui/litellm-dashboard/src/utils/roles.ts +++ b/ui/litellm-dashboard/src/utils/roles.ts @@ -28,6 +28,12 @@ export const isAdminRole = (role: string): boolean => { return all_admin_roles.includes(role); }; +// /user/list admits proxy admins and org admins; the session role for the latter is the formatted +// "Org Admin", which all_admin_roles does not carry +const rolesAllowedToListUsers: string[] = [...all_admin_roles, "Org Admin"]; + +export const canListUsers = (role: string | null): boolean => rolesAllowedToListUsers.includes(role ?? ""); + export const isProxyAdminRole = (role: string): boolean => { return role === "proxy_admin" || role === "Admin"; }; From b7596d6fba97f712a2dca3ead9ed280af3654292 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:16:30 +0000 Subject: [PATCH 005/102] refactor(ui): drop redundant comment on canListUsers role list Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/utils/roles.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/utils/roles.ts b/ui/litellm-dashboard/src/utils/roles.ts index 1cb7e75c19b..066a83992b4 100644 --- a/ui/litellm-dashboard/src/utils/roles.ts +++ b/ui/litellm-dashboard/src/utils/roles.ts @@ -28,8 +28,6 @@ export const isAdminRole = (role: string): boolean => { return all_admin_roles.includes(role); }; -// /user/list admits proxy admins and org admins; the session role for the latter is the formatted -// "Org Admin", which all_admin_roles does not carry const rolesAllowedToListUsers: string[] = [...all_admin_roles, "Org Admin"]; export const canListUsers = (role: string | null): boolean => rolesAllowedToListUsers.includes(role ?? ""); From 97dbd2dfbf3de0386efbe77057bbf75bc1aa1336 Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 22:16:19 +0530 Subject: [PATCH 006/102] fix(gemini): preserve candidates with finishReason and no content (#40477) --- litellm/litellm_core_utils/core_helpers.py | 1 + .../adapters/transformation.py | 2 + .../vertex_and_google_ai_studio_gemini.py | 30 ++-- .../transformation.py | 20 ++- ...test_vertex_and_google_ai_studio_gemini.py | 146 ++++++++++++++++++ 5 files changed, 184 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index aa7d6ca1699..66180b165f8 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -224,6 +224,7 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { "IMAGE_PROHIBITED_CONTENT": "content_filter", "TOO_MANY_TOOL_CALLS": "stop", "MALFORMED_RESPONSE": "stop", + "NO_IMAGE": "content_filter", # Zhipu GLM "network_error": "stop", "sensitive": "content_filter", diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 8ff9f2e0679..f4cc569bcef 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1367,6 +1367,8 @@ class LiteLLMAnthropicMessagesAdapter: return "max_tokens" elif openai_finish_reason == "tool_calls": return "tool_use" + elif openai_finish_reason in ["content_filter", "refusal"]: + return "refusal" return "end_turn" @staticmethod diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d113b2b4f6b..01d1f063b1b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1340,6 +1340,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT", "TOO_MANY_TOOL_CALLS", "MALFORMED_RESPONSE", + "NO_IMAGE", } ) @@ -2224,22 +2225,23 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): grounding_metadata: Final[list[dict]] = [] url_context_metadata: Final[list[dict]] = [] - image_response: list[ImageURLListItem] | None = None safety_ratings: Final[list] = [] citation_metadata: Final[list] = [] - chat_completion_message: Final[ChatCompletionResponseMessage] = {"role": "assistant"} - chat_completion_logprobs: ChoiceLogprobs | None = None - tools: list[ChatCompletionToolCallChunk] | None = [] - functions: ChatCompletionToolCallFunctionChunk | None = None - thinking_blocks: list[ChatCompletionThinkingBlock] | None = None - reasoning_content: str | None = None - thought_signatures: Sequence[str] | None = None - server_side_tool_invocations: list[dict[str, object]] | None = None for idx, candidate in enumerate(_candidates): - if "content" not in candidate: + if "content" not in candidate and "finishReason" not in candidate: continue + image_response: list[ImageURLListItem] | None = None + chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} + chat_completion_logprobs: ChoiceLogprobs | None = None + tools: list[ChatCompletionToolCallChunk] | None = [] + functions: ChatCompletionToolCallFunctionChunk | None = None + thinking_blocks: list[ChatCompletionThinkingBlock] | None = None + reasoning_content: str | None = None + thought_signatures: Sequence[str] | None = None + server_side_tool_invocations: list[dict[str, object]] | None = None + # Extract metadata using helper function ( candidate_grounding_metadata, @@ -2253,7 +2255,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): safety_ratings.extend(candidate_safety_ratings) citation_metadata.extend(candidate_citation_metadata) - if "parts" in candidate["content"]: + if "content" in candidate and candidate["content"] and "parts" in candidate["content"]: ( content, reasoning_content, @@ -2348,6 +2350,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_invocation_fields["server_side_tool_invocations"] = server_side_tool_invocations chat_completion_message["provider_specific_fields"] = tool_invocation_fields + if candidate.get("finishReason"): + finish_reason_fields = chat_completion_message.get("provider_specific_fields") or {} + finish_reason_fields["native_finish_reason"] = candidate.get("finishReason") + chat_completion_message["provider_specific_fields"] = finish_reason_fields + if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( chat_completion_message=chat_completion_message, @@ -2368,6 +2375,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): message=chat_completion_message, logprobs=chat_completion_logprobs, enhancements=None, + provider_specific_fields=chat_completion_message.get("provider_specific_fields"), ) model_response.choices.append(choice) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fca5b0d11cf..13119085e46 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2272,13 +2272,27 @@ class LiteLLMCompletionResponsesConfig: if choices and len(choices) > 0: finish_reason = choices[0].finish_reason + status: Final[ResponsesAPIStatus] = ( + LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( + finish_reason + ) + ) + incomplete_details = getattr(chat_completion_response, "incomplete_details", None) + if incomplete_details is None and status == "incomplete": + from openai.types.responses.response import IncompleteDetails + + if finish_reason == "length": + incomplete_details = IncompleteDetails(reason="max_output_tokens") + elif finish_reason in ["content_filter", "refusal"]: + incomplete_details = IncompleteDetails(reason="content_filter") + responses_api_response: Final[ResponsesAPIResponse] = ResponsesAPIResponse( id=chat_completion_response.id, created_at=chat_completion_response.created, model=chat_completion_response.model, object="response", error=getattr(chat_completion_response, "error", None), - incomplete_details=getattr(chat_completion_response, "incomplete_details", None), + incomplete_details=incomplete_details, instructions=getattr(chat_completion_response, "instructions", None), metadata=getattr(chat_completion_response, "metadata", {}), output=LiteLLMCompletionResponsesConfig._transform_chat_completion_choices_to_responses_output( @@ -2296,9 +2310,7 @@ class LiteLLMCompletionResponsesConfig: max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), previous_response_id=getattr(chat_completion_response, "previous_response_id", None), reasoning=None, - status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( - finish_reason - ), + status=status, text={}, truncation=getattr(chat_completion_response, "truncation", None), usage=LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 101f6e6fa5d..a048ad4f171 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5836,3 +5836,149 @@ def test_supported_reasoning_efforts_still_map(model): drop_params=False, ) assert "thinkingConfig" in result + + +def test_gemini_candidate_with_finish_reason_no_content_chat_completion(): + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + model_response = ModelResponse() + logging_obj = MagicMock() + raw_response = MagicMock() + raw_response.headers = {} + + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=model_response, + model="gemini-2.5-flash-image", + logging_obj=logging_obj, + raw_response=raw_response, + ) + assert len(resp.choices) == 1 + assert resp.choices[0].finish_reason == "content_filter" + assert resp.choices[0].message.content is None + assert resp.choices[0].provider_specific_fields["native_finish_reason"] == "NO_IMAGE" + + +def test_gemini_candidate_with_finish_reason_no_content_anthropic_messages(): + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash-image", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_resp = adapter.translate_openai_response_to_anthropic( + response=resp, + tool_name_mapping={}, + ) + assert anthropic_resp["stop_reason"] == "refusal" + assert anthropic_resp["content"] == [] + + +def test_gemini_candidate_with_finish_reason_no_content_responses_api(): + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash-image", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + responses_resp = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Generate picture", + responses_api_request={}, + chat_completion_response=resp, + ) + assert responses_resp.status == "incomplete" + assert responses_resp.incomplete_details is not None + assert responses_resp.incomplete_details.reason == "content_filter" + + +def test_gemini_candidate_other_finish_reasons_no_content(): + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + config = VertexGeminiConfig() + max_tokens_response = { + "candidates": [{"finishReason": "MAX_TOKENS", "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 50, "totalTokenCount": 60}, + } + resp_length = config._transform_google_generate_content_to_openai_model_response( + completion_response=max_tokens_response, + model_response=ModelResponse(), + model="gemini-2.5-flash", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + assert len(resp_length.choices) == 1 + assert resp_length.choices[0].finish_reason == "length" + assert resp_length.choices[0].provider_specific_fields["native_finish_reason"] == "MAX_TOKENS" + + anthropic_length = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( + response=resp_length, + tool_name_mapping={}, + ) + assert anthropic_length["stop_reason"] == "max_tokens" + + responses_length = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="thinking request", + responses_api_request={}, + chat_completion_response=resp_length, + ) + assert responses_length.status == "incomplete" + assert responses_length.incomplete_details.reason == "max_output_tokens" + From cd66b34b45dace036126f4ea809df132407d907c Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 22:48:08 +0530 Subject: [PATCH 007/102] style(responses): apply ruff formatting to transformation.py --- .../litellm_completion_transformation/transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 13119085e46..10e56e85bff 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2273,9 +2273,7 @@ class LiteLLMCompletionResponsesConfig: finish_reason = choices[0].finish_reason status: Final[ResponsesAPIStatus] = ( - LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( - finish_reason - ) + LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status(finish_reason) ) incomplete_details = getattr(chat_completion_response, "incomplete_details", None) if incomplete_details is None and status == "incomplete": From e54399eff7a13e649b6a353486922166b1bf5688 Mon Sep 17 00:00:00 2001 From: tusharjamunkar Date: Sat, 12 Sep 2026 23:07:31 +0530 Subject: [PATCH 008/102] test: add direct coverage for content_filter and refusal in anthropic and responses adapters --- ...al_pass_through_adapters_transformation.py | 40 ++++++++++++ .../test_litellm_completion_responses.py | 62 +++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 03b9840b1c3..0e09e27f4db 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -102,6 +102,46 @@ def test_translate_chat_length_takes_precedence_over_refusal(): assert result.get("stop_details") is None +def test_translate_chat_content_filter_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-content-filter", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="content_filter", + message=Message(content=None, role="assistant"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "refusal" + + +def test_translate_chat_refusal_finish_reason_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-refusal-reason", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="refusal", + message=Message(content=None, role="assistant"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "refusal" + + def test_translate_streaming_openai_chunk_to_anthropic_content_block(): choices = [ StreamingChoices( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 46249e50572..f2ebbf316c3 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4246,3 +4246,65 @@ class TestStreamingSnapshotItemIds: reasoning_items = _bridged_output_items(completed_event.response, "reasoning") assert len(reasoning_items) == 1 assert reasoning_items[0].id == streamed_event.item_id + + +def test_transform_chat_completion_response_incomplete_details(): + from openai.types.responses.response import IncompleteDetails + + resp_length = ModelResponse( + id="resp-length", + choices=[Choices(index=0, finish_reason="length", message=Message(content="cutoff", role="assistant"))], + model="gpt-4o", + ) + result_length = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_length, + ) + assert result_length.status == "incomplete" + assert result_length.incomplete_details is not None + assert result_length.incomplete_details.reason == "max_output_tokens" + + resp_filter = ModelResponse( + id="resp-filter", + choices=[Choices(index=0, finish_reason="content_filter", message=Message(content=None, role="assistant"))], + model="gpt-4o", + ) + result_filter = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_filter, + ) + assert result_filter.status == "incomplete" + assert result_filter.incomplete_details is not None + assert result_filter.incomplete_details.reason == "content_filter" + + resp_refusal = ModelResponse( + id="resp-refusal", + choices=[Choices(index=0, finish_reason="refusal", message=Message(content=None, role="assistant"))], + model="gpt-4o", + ) + result_refusal = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_refusal, + ) + assert result_refusal.status == "incomplete" + assert result_refusal.incomplete_details is not None + assert result_refusal.incomplete_details.reason == "content_filter" + + existing_details = IncompleteDetails(reason="content_filter") + resp_existing = ModelResponse( + id="resp-existing", + choices=[Choices(index=0, finish_reason="length", message=Message(content="cutoff", role="assistant"))], + model="gpt-4o", + ) + resp_existing.incomplete_details = existing_details + result_existing = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_existing, + ) + assert result_existing.status == "incomplete" + assert result_existing.incomplete_details == existing_details + From cd8887d72cef4581519914d99be1da153fbe8ee8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 09:50:11 +0000 Subject: [PATCH 009/102] fix(mistral): accept reasoning_effort on all models and drop client_metadata for Codex compatibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/mistral/chat/transformation.py | 14 +++++--- .../test_mistral_chat_transformation.py | 33 +++++++++++++++++-- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index a76a8a3e98c..50aefcdc918 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -99,11 +99,11 @@ class MistralConfig(OpenAIGPTConfig): "stop", "response_format", "parallel_tool_calls", + "reasoning_effort", ] - # Add reasoning support for magistral models if "magistral" in model.lower(): - supported_params.extend(["thinking", "reasoning_effort"]) + supported_params.append("thinking") return supported_params @@ -171,9 +171,11 @@ class MistralConfig(OpenAIGPTConfig): optional_params["extra_body"] = {"random_seed": value} if param == "response_format": optional_params["response_format"] = value - if param == "reasoning_effort" and "magistral" in model.lower(): - # Flag that we need to add reasoning system prompt - optional_params["_add_reasoning_prompt"] = True + if param == "reasoning_effort": + if "magistral" in model.lower(): + optional_params["_add_reasoning_prompt"] = True + else: + optional_params["reasoning_effort"] = value if param == "thinking" and "magistral" in model.lower(): # Flag that we need to add reasoning system prompt optional_params["_add_reasoning_prompt"] = True @@ -534,6 +536,8 @@ class MistralConfig(OpenAIGPTConfig): if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False): messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params) + optional_params.pop("client_metadata", None) + # Call parent transform_request which handles _transform_messages return super().transform_request( model=model, diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py index 15694d9f218..edfaf352e1f 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py @@ -51,11 +51,11 @@ class TestMistralReasoningSupport: assert "reasoning_effort" in supported_params assert "thinking" in supported_params - # Test non-magistral model doesn't include reasoning parameters + # Non-magistral models accept reasoning_effort (forwarded verbatim) but not thinking supported_params_normal = mistral_config.get_supported_openai_params( "mistral/mistral-large-latest" ) - assert "reasoning_effort" not in supported_params_normal + assert "reasoning_effort" in supported_params_normal assert "thinking" not in supported_params_normal def test_map_openai_params_reasoning_effort(self): @@ -73,7 +73,7 @@ class TestMistralReasoningSupport: assert result.get("_add_reasoning_prompt") is True - # Test reasoning_effort ignored for non-magistral model + # Test reasoning_effort forwarded verbatim for non-magistral model optional_params_normal = {} result_normal = mistral_config.map_openai_params( non_default_params={"reasoning_effort": "low"}, @@ -83,6 +83,33 @@ class TestMistralReasoningSupport: ) assert "_add_reasoning_prompt" not in result_normal + assert result_normal["reasoning_effort"] == "low" + + def test_reasoning_effort_not_unsupported_for_non_magistral(self): + """Codex sends reasoning_effort to every model; Mistral must not raise UnsupportedParamsError.""" + import litellm + + optional_params = litellm.get_optional_params( + model="mistral-medium-latest", + custom_llm_provider="mistral", + reasoning_effort="medium", + ) + assert optional_params["reasoning_effort"] == "medium" + + def test_client_metadata_stripped_from_request(self): + """client_metadata passed by Codex must not reach Mistral, whose schema rejects unknown fields.""" + mistral_config = MistralConfig() + + request = mistral_config.transform_request( + model="mistral-medium-latest", + messages=[{"role": "user", "content": "hi"}], + optional_params={"client_metadata": {"originator": "codex_cli_rs"}, "temperature": 0.2}, + litellm_params={}, + headers={}, + ) + + assert "client_metadata" not in request + assert request["temperature"] == 0.2 def test_map_openai_params_thinking(self): """Test that thinking parameter is properly mapped for magistral models.""" From a8305129a7ff0b8389411b27935008536d698aec Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 09:56:08 +0000 Subject: [PATCH 010/102] refactor(mistral): keep map_openai_params under the complexity ceiling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/mistral/chat/transformation.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 50aefcdc918..970da0582ae 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -171,12 +171,9 @@ class MistralConfig(OpenAIGPTConfig): optional_params["extra_body"] = {"random_seed": value} if param == "response_format": optional_params["response_format"] = value - if param == "reasoning_effort": - if "magistral" in model.lower(): - optional_params["_add_reasoning_prompt"] = True - else: - optional_params["reasoning_effort"] = value - if param == "thinking" and "magistral" in model.lower(): + if param == "reasoning_effort" and "magistral" not in model.lower(): + optional_params["reasoning_effort"] = value + if param in ("reasoning_effort", "thinking") and "magistral" in model.lower(): # Flag that we need to add reasoning system prompt optional_params["_add_reasoning_prompt"] = True if param == "parallel_tool_calls": From dd209ba97b3730523b0a3b3a0c00e84acbb89a9a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:25:53 +0000 Subject: [PATCH 011/102] fix(bedrock): carry s3_endpoint_url and s3_region_name into file content downloads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 2 ++ litellm/litellm_core_utils/get_litellm_params.py | 2 ++ litellm/types/utils.py | 1 + tests/test_litellm/batches/test_batch_utils.py | 2 ++ .../litellm_core_utils/test_get_litellm_params.py | 12 ++++++++++++ .../files/test_bedrock_files_transformation.py | 13 +++++++++++++ 6 files changed, 32 insertions(+) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 87f8fd3946e..0ec39ebf2b2 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -530,6 +530,8 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: "vertex_credentials", "gcs_bucket_name", "bucket_name", + "s3_endpoint_url", + "s3_region_name", "timeout", "max_retries", "_litellm_internal_model_credentials", diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index edd2e88f95c..a70dce89680 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -43,6 +43,8 @@ OPTIONAL_KWARGS_KEYS: Final = ( "timeout", "gcs_bucket_name", "bucket_name", + "s3_endpoint_url", + "s3_region_name", "vertex_credentials", "vertex_project", "vertex_location", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 1d73542c9bb..ee8a3956be8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3721,6 +3721,7 @@ bedrock_batch_litellm_params: Final = ( "aws_batch_role_arn", "s3_bucket_name", "s3_region_name", + "s3_endpoint_url", "s3_output_bucket_name", "bedrock_tags", ) diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 768ea332677..8c8e0621b07 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -278,6 +278,8 @@ def test_extract_credentials_all_supported_keys(): "vertex_credentials", "gcs_bucket_name", "bucket_name", + "s3_endpoint_url", + "s3_region_name", "timeout", "max_retries", } diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index f026ff57719..a34bc2af59d 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -55,6 +55,18 @@ class TestGetLitellmParamsKwargsExtraction: assert result["timeout"] == 30 assert result["rpm"] == 100 + def test_s3_endpoint_kwargs_are_extracted_when_provided(self): + result = get_litellm_params( + s3_endpoint_url="https://bucket.vpce-abc.s3.us-east-1.vpce.amazonaws.com", + s3_region_name="us-east-1", + ) + assert result["s3_endpoint_url"] == "https://bucket.vpce-abc.s3.us-east-1.vpce.amazonaws.com" + assert result["s3_region_name"] == "us-east-1" + + result_without_s3_kwargs = get_litellm_params() + assert "s3_endpoint_url" not in result_without_s3_kwargs + assert "s3_region_name" not in result_without_s3_kwargs + def test_subset_of_kwargs_only_includes_provided(self): """Only provided kwargs appear, others remain absent.""" result = get_litellm_params(azure_ad_token="token123") diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c609455f3d8..d5bfe5cdfc7 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -2271,6 +2271,19 @@ class TestBedrockFileContentTransformation: authorization = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]["Authorization"] assert "/eu-west-1/s3/aws4_request" in authorization + def test_s3_request_target_uses_configured_endpoint_url(self): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + lp = get_litellm_params( + aws_region_name="us-east-1", + s3_endpoint_url="https://bucket.vpce-abc.s3.us-east-1.vpce.amazonaws.com", + ) + + assert BedrockFilesConfig()._s3_request_target( + optional_params={}, litellm_params=lp + ).endpoint_url == "https://bucket.vpce-abc.s3.us-east-1.vpce.amazonaws.com" + def test_validate_environment_merges_and_pops_signed_get_headers(self): from litellm.llms.bedrock.files.transformation import ( S3_SIGNED_REQUEST_HEADERS_PARAM, From 817c396383e56db8055b9b5601771107ca201bf7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:39:09 +0000 Subject: [PATCH 012/102] fix(bedrock): preserve S3 endpoint in credential snapshots Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/router.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/types/router.py b/litellm/types/router.py index 0aefc07ae4b..1ce86479f34 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -299,6 +299,7 @@ class CredentialLiteLLMParams(BaseModel): aws_bedrock_runtime_endpoint: str | None = None aws_bedrock_project_id: str | None = None s3_bucket_name: str | None = None + s3_endpoint_url: str | None = None s3_region_name: str | None = None s3_encryption_key_id: str | None = None aws_batch_role_arn: str | None = None From 3b620c65d25ea15ed5d955ae6b31dbb697fa789f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:43:27 +0000 Subject: [PATCH 013/102] chore: sync schema.d.ts with proxy OpenAPI spec Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0b0e3e18215..6e471d42e49 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29899,6 +29899,8 @@ export interface components { s3_bucket_name?: string | null; /** S3 Encryption Key Id */ s3_encryption_key_id?: string | null; + /** S3 Endpoint Url */ + s3_endpoint_url?: string | null; /** S3 Output Bucket Name */ s3_output_bucket_name?: string | null; /** S3 Region Name */ @@ -40113,6 +40115,8 @@ export interface components { s3_bucket_name?: string | null; /** S3 Encryption Key Id */ s3_encryption_key_id?: string | null; + /** S3 Endpoint Url */ + s3_endpoint_url?: string | null; /** S3 Output Bucket Name */ s3_output_bucket_name?: string | null; /** S3 Region Name */ From 911f66aff69fabb6666bde3f54db70960cb04b56 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:12:57 -0500 Subject: [PATCH 014/102] feat(azure_ai): support FLUX.2 flex images --- litellm/images/main.py | 9 +- litellm/images/utils.py | 7 +- .../litellm_core_utils/llm_cost_calc/utils.py | 3 + .../image_edit/flux2_transformation.py | 64 ++++--- .../image_generation/cost_calculator.py | 34 +++- .../image_generation/flux_transformation.py | 91 +++++++-- ...odel_prices_and_context_window_backup.json | 19 ++ litellm/proxy/_lazy_openapi_snapshot.json | 2 +- litellm/types/llms/openai.py | 4 + model_prices_and_context_window.json | 19 ++ ...test_azure_ai_image_edit_transformation.py | 116 ++++++++++++ .../test_azure_ai_flux2_image_generation.py | 172 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 13 files changed, 491 insertions(+), 53 deletions(-) create mode 100644 tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 6a94e7c8df2..81547a153c3 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -846,7 +846,12 @@ def image_edit( local_vars.update(kwargs) # Get ImageEditOptionalRequestParams with only valid parameters image_edit_optional_params: Final[ImageEditOptionalRequestParams] = ( - _get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars) + _get_ImageEditRequestUtils().get_requested_image_edit_optional_param( + local_vars, + provider_supported_params=frozenset( + image_edit_provider_config.get_supported_openai_params(model) + ).intersection(non_default_params), + ) ) # Get optional parameters for the responses API image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit( @@ -857,7 +862,7 @@ def image_edit( additional_drop_params=kwargs.get("additional_drop_params"), ) - if ( + if image_edit_provider_config.use_multipart_form_data() and ( custom_llm_provider == "openai" or custom_llm_provider == "azure" or custom_llm_provider in litellm.openai_compatible_providers diff --git a/litellm/images/utils.py b/litellm/images/utils.py index 49b70870de6..24454954714 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Collection, Mapping from io import BufferedReader, BytesIO from typing import Any, Final, cast, get_type_hints @@ -63,6 +63,7 @@ class ImageEditRequestUtils: @staticmethod def get_requested_image_edit_optional_param( params: Mapping[str, object], + provider_supported_params: Collection[str] = (), ) -> ImageEditOptionalRequestParams: """ Filter parameters to only include those defined in ImageEditOptionalRequestParams. @@ -73,7 +74,9 @@ class ImageEditRequestUtils: Returns: ImageEditOptionalRequestParams instance with only the valid parameters """ - valid_keys: Final = get_type_hints(ImageEditOptionalRequestParams).keys() + valid_keys: Final = frozenset(get_type_hints(ImageEditOptionalRequestParams)) | frozenset( + provider_supported_params + ) filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None} return cast(ImageEditOptionalRequestParams, filtered_params) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index baa9aab1087..0a0e92ff3a3 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1853,6 +1853,9 @@ class CostCalculatorUtils: return azure_ai_image_cost_calculator( model=model, image_response=completion_response, + size=resolved_size, + n=resolved_n, + optional_params=optional_params, ) elif custom_llm_provider == litellm.LlmProviders.FAL_AI.value: from litellm.llms.fal_ai.cost_calculator import ( diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index a09a80985b7..aa8905e5601 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -1,5 +1,7 @@ import base64 +from collections.abc import Mapping, Sequence from io import BufferedReader +from types import MappingProxyType from typing import Any, Final from httpx._types import RequestFiles @@ -24,7 +26,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): Azure AI Foundry FLUX 2 image edit config Supports FLUX 2 models (e.g., flux.2-pro) for image editing. - Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation, + Uses the model-specific /providers/blackforestlabs/v1/flux-2-* endpoint as image generation, with the image passed as base64 in JSON body. """ @@ -33,11 +35,17 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): FLUX 2 supports a subset of OpenAI image edit params """ return [ - "prompt", - "image", - "model", "n", "size", + "width", + "height", + "num_images", + "seed", + "safety_tolerance", + "output_format", + "aspect_ratio", + "guidance", + "steps", ] def map_openai_params( @@ -50,14 +58,14 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): Map OpenAI params to FLUX 2 params. FLUX 2 uses the same param names as OpenAI for supported params. """ - mapped_params: Final[dict[str, Any]] = {} - supported_params: Final = self.get_supported_openai_params(model) - - for key, value in dict(image_edit_optional_params).items(): - if key in supported_params and value is not None: - mapped_params[key] = value - - return mapped_params + return AzureFoundryFluxImageGenerationConfig().map_openai_params( + non_default_params=MappingProxyType( + {key: value for key, value in image_edit_optional_params.items() if value is not None} + ), + optional_params=MappingProxyType({}), + model=model, + drop_params=drop_params, + ) def use_multipart_form_data(self) -> bool: """FLUX 2 uses JSON requests, not multipart/form-data.""" @@ -90,7 +98,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): self, model: str, prompt: str | None, - image: FileTypes | None, + image: FileTypes | Sequence[FileTypes] | None, image_edit_optional_request_params: dict, litellm_params: GenericLiteLLMParams, headers: dict, @@ -107,29 +115,29 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): if image is None: raise ValueError("FLUX 2 image edit requires an image.") - image_b64: Final = self._convert_image_to_base64(image) + images: Final = tuple(image) if isinstance(image, list) else (image,) + if not images: + raise ValueError("FLUX 2 image edit requires at least one image.") + max_reference_images: Final = 10 if "flex" in model.lower() else 8 + if len(images) > max_reference_images: + raise ValueError(f"{model} supports at most {max_reference_images} reference images.") - # Build request body with required params + reference_images: Final[Mapping[str, str]] = MappingProxyType( + { + "input_image" if index == 1 else f"input_image_{index}": self._convert_image_to_base64(reference_image) + for index, reference_image in enumerate(images, start=1) + } + ) request_body: Final[dict[str, Any]] = { "prompt": prompt, - "image": image_b64, "model": model, + **reference_images, + **image_edit_optional_request_params, } - - # Add mapped optional params (already filtered by map_openai_params) - request_body.update(image_edit_optional_request_params) - - # Return JSON body and empty files list (FLUX 2 doesn't use multipart) return request_body, [] def _convert_image_to_base64(self, image: Any) -> str: """Convert image file to base64 string""" - # Handle list of images (take first one) - if isinstance(image, list): - if len(image) == 0: - raise ValueError("Empty image list provided") - image = image[0] - if isinstance(image, BufferedReader): image_bytes = image.read() image.seek(0) # Reset file pointer for potential reuse @@ -151,7 +159,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ Constructs a complete URL for Azure AI Foundry FLUX 2 image edits. - Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation. + Uses the same model-specific BFL provider endpoint as image generation. """ api_base = AzureFoundryModelInfo.get_api_base(api_base) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 106c7e42b83..086293f26c0 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final import litellm @@ -10,6 +11,9 @@ from litellm.types.utils import ImageResponse def cost_calculator( model: str, image_response: Any, + size: str | None = None, + n: int | None = None, + optional_params: Mapping[str, object] | None = None, ) -> float: """ Azure AI image generation cost calculator @@ -28,10 +32,32 @@ def cost_calculator( if token_based_cost is not None: return token_based_cost + num_images: Final = n if n is not None else len(image_response.data or ()) output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0 - num_images: int = 0 - if image_response.data: - num_images = len(image_response.data) - return output_cost_per_image * num_images + if output_cost_per_image: + return output_cost_per_image * num_images + + model_cost: Final = litellm.model_cost[_model_info["key"]] + input_cost_per_pixel: Final[float] = model_cost.get("input_cost_per_pixel") or 0.0 + if input_cost_per_pixel: + from litellm.cost_calculator import default_image_cost_calculator + + cost_model: Final = ( + model if model.startswith(f"{litellm.LlmProviders.AZURE_AI.value}/") else f"azure_ai/{model}" + ) + width: Final = optional_params.get("width") if optional_params else None + height: Final = optional_params.get("height") if optional_params else None + pixel_size: Final = ( + f"{width}x{height}" + if type(width) is int and type(height) is int and width > 0 and height > 0 + else size or image_response.size + ) + return default_image_cost_calculator( + model=cost_model, + custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, + size=pixel_size, + n=num_images, + ) + return 0.0 raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 65b5a35af52..b10e8a6f35e 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -1,18 +1,13 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from litellm.llms.openai.image_generation import GPTImageGenerationConfig +from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): - """ - Azure Foundry flux image generation config - - From manual testing it follows the gpt-image-1 image generation config - - (Azure Foundry does not have any docs on supported params at the time of writing) - - From our test suite - following GPTImageGenerationConfig is working for this model - """ + """Azure Foundry BFL API configuration for FLUX image generation.""" @staticmethod def get_flux2_image_generation_url( @@ -25,11 +20,11 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): FLUX 2 models on Azure AI use a different URL pattern than standard Azure OpenAI: - Standard: /openai/deployments/{model}/images/generations - - FLUX 2: /providers/blackforestlabs/v1/flux-2-pro + - FLUX 2: /providers/blackforestlabs/v1/{model-path} Args: api_base: Base URL (e.g., https://litellm-ci-cd-prod.services.ai.azure.com) - model: Model name (e.g., flux.2-pro) + model: Model name (e.g., FLUX.2-flex or FLUX.2-pro) api_version: API version (e.g., preview) Returns: @@ -47,9 +42,8 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): return api_base return f"{api_base}?api-version={api_version}" - # Construct the FLUX 2 provider path - # Model name flux.2-pro maps to endpoint flux-2-pro - return f"{api_base}/providers/blackforestlabs/v1/flux-2-pro?api-version={api_version}" + provider_model_path: Final = AzureFoundryFluxImageGenerationConfig.get_flux2_provider_model_path(model) + return f"{api_base}/providers/blackforestlabs/v1/{provider_model_path}?api-version={api_version}" @staticmethod def is_flux2_model(model: str) -> bool: @@ -64,3 +58,72 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): """ model_lower: Final = model.lower().replace(".", "-").replace("_", "-") return "flux-2" in model_lower or "flux2" in model_lower + + @staticmethod + def get_flux2_provider_model_path(model: str) -> str: + normalized_model: Final = model.lower().replace(".", "-").replace("_", "-") + return "flux-2-flex" if "flux-2-flex" in normalized_model else "flux-2-pro" + + def get_supported_openai_params( # mutable-ok: inherited config contract returns a list + self, model: str + ) -> list[OpenAIImageGenerationOptionalParams]: + if not self.is_flux2_model(model): + return super().get_supported_openai_params(model) + return [ # mutable-ok: BaseImageGenerationConfig requires a list + "n", + "size", + "output_format", + "seed", + "safety_tolerance", + "aspect_ratio", + "width", + "height", + "num_images", + "guidance", + "steps", + ] + + @staticmethod + def _map_parameter(name: str, value: object) -> tuple[tuple[str, object], ...]: + if name == "n": + return (("num_images", value),) + if name != "size": + return ((name, value),) + + try: + width, height = (int(dimension) for dimension in str(value).lower().split("x")) + except (TypeError, ValueError): + raise ValueError(f"Invalid size format '{value}'. Expected 'WxH', for example '1024x1024'.") + return (("width", width), ("height", height)) + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: inherited config contract returns a dict + if not self.is_flux2_model(model): + return super().map_openai_params( + non_default_params=dict(non_default_params), + optional_params=dict(optional_params), + model=model, + drop_params=drop_params, + ) + supported_params: Final = self.get_supported_openai_params(model) + unsupported_params: Final = tuple(name for name in non_default_params if name not in supported_params) + if unsupported_params and not drop_params: + raise ValueError( + f"Parameters {unsupported_params} are not supported for model {model}. " + f"Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." + ) + + mapped_params: Final[Mapping[str, object]] = MappingProxyType( + { + mapped_name: mapped_value + for name, value in non_default_params.items() + if name in supported_params + for mapped_name, mapped_value in self._map_parameter(name, value) + } + ) + return {**optional_params, **mapped_params} # mutable-ok: inherited config contract returns a dict diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9f91cf82f41..bcb2dfc53e8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9900,6 +9900,25 @@ "/v1/images/generations" ] }, + "azure_ai/FLUX.2-flex": { + "input_cost_per_pixel": 5e-08, + "litellm_provider": "azure_ai", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "image_generation", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "image" + ] + }, "azure_ai/FW-DeepSeek-V3.2": { "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.1e-07, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index fb2f014e3d8..81889e4da0c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -18974,7 +18974,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e3eac9b9205..710b34116e5 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1160,6 +1160,10 @@ OpenAIImageGenerationOptionalParams = Literal[ "image_url", "image_prompt_strength", "aspect_ratio", + "width", + "height", + "guidance", + "steps", "imageConfig", ] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9f91cf82f41..bcb2dfc53e8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9900,6 +9900,25 @@ "/v1/images/generations" ] }, + "azure_ai/FLUX.2-flex": { + "input_cost_per_pixel": 5e-08, + "litellm_provider": "azure_ai", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "image_generation", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "image" + ] + }, "azure_ai/FW-DeepSeek-V3.2": { "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.1e-07, diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index b6cb7ea9b54..94b23c5f1c2 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -1,12 +1,20 @@ +import base64 +import json +from collections.abc import Mapping +from typing import Final +import httpx +import pytest import litellm +from litellm.images.utils import ImageEditRequestUtils from litellm.llms.azure_ai.image_edit.flux2_transformation import ( AzureFoundryFlux2ImageEditConfig, ) from litellm.llms.azure_ai.image_edit.transformation import ( AzureFoundryFluxImageEditConfig, ) +from litellm.llms.custom_httpx.http_handler import HTTPHandler def test_azure_ai_validate_environment(): @@ -60,3 +68,111 @@ def test_flux2_validate_environment_with_entra_token(monkeypatch): assert headers["Authorization"] == "Bearer entra-token" assert headers["Content-Type"] == "application/json" + + +def test_flux2_image_edit_maps_openai_and_provider_parameters(): + config = AzureFoundryFlux2ImageEditConfig() + requested_params = ImageEditRequestUtils.get_requested_image_edit_optional_param( + { + "n": 2, + "size": "1536x1024", + "guidance": 4.5, + "steps": 32, + "unrelated": "discarded", + }, + provider_supported_params=config.get_supported_openai_params("FLUX.2-flex"), + ) + mapped_params = config.map_openai_params( + image_edit_optional_params=requested_params, + model="FLUX.2-flex", + drop_params=False, + ) + + assert mapped_params == { + "num_images": 2, + "width": 1536, + "height": 1024, + "guidance": 4.5, + "steps": 32, + } + + +@pytest.mark.parametrize( + ("model", "max_reference_images"), + [ + ("FLUX.2-flex", 10), + ("FLUX.2-pro", 8), + ], +) +def test_flux2_image_edit_uses_all_reference_fields(model: str, max_reference_images: int): + images = [f"image-{index}".encode() for index in range(1, max_reference_images + 1)] + request, files = AzureFoundryFlux2ImageEditConfig().transform_image_edit_request( + model=model, + prompt="Blend every reference", + image=images, + image_edit_optional_request_params={"guidance": 4.5, "steps": 20}, + litellm_params={}, + headers={}, + ) + + assert files == [] + assert request["input_image"] == base64.b64encode(images[0]).decode() + assert request[f"input_image_{max_reference_images}"] == base64.b64encode(images[-1]).decode() + assert "input_image_1" not in request + assert "image" not in request + assert len([key for key in request if key.startswith("input_image")]) == max_reference_images + assert request["guidance"] == 4.5 + assert request["steps"] == 20 + + +@pytest.mark.parametrize( + ("model", "reference_images"), + [ + ("FLUX.2-flex", 11), + ("FLUX.2-pro", 9), + ], +) +def test_flux2_image_edit_rejects_too_many_references(model: str, reference_images: int): + with pytest.raises(ValueError, match=f"at most {reference_images - 1} reference images"): + AzureFoundryFlux2ImageEditConfig().transform_image_edit_request( + model=model, + prompt="Blend every reference", + image=[b"image"] * reference_images, + image_edit_optional_request_params={}, + litellm_params={}, + headers={}, + ) + + +@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024})) +@pytest.mark.usefixtures("local_model_cost_map") +def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[str, int | str]): + def respond(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + assert body == { + "model": "FLUX.2-flex", + "prompt": "Add a hat", + "input_image": base64.b64encode(b"image").decode(), + "num_images": 2, + "width": 2048, + "height": 1024, + "guidance": 4.5, + "steps": 32, + } + return httpx.Response(200, json={"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]}) + + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + response: Final = litellm.image_edit( + model="azure_ai/FLUX.2-flex", + image=b"image", + prompt="Add a hat", + api_key="test-key", + api_base="https://example.services.ai.azure.com", + client=client, + n=2, + guidance=4.5, + steps=32, + **dimensions, + ) + + assert response._hidden_params["response_cost"] == pytest.approx(5e-08 * 2048 * 1024 * 2) diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py new file mode 100644 index 00000000000..07026cf4309 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -0,0 +1,172 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils +from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.azure.image_generation import get_azure_image_generation_config +from litellm.llms.azure.image_generation.http_utils import azure_deployment_image_generation_json_body +from litellm.llms.azure_ai.image_generation.flux_transformation import ( + AzureFoundryFluxImageGenerationConfig, +) +from litellm.types.utils import ImageObject, ImageResponse +from litellm.utils import _invalidate_model_cost_lowercase_map + + +@pytest.fixture(autouse=True) +def use_local_model_cost_map(monkeypatch): + monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) + litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() + yield + litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize( + ("model", "provider_path"), + [ + ("FLUX.2-flex", "flux-2-flex"), + ("FLUX.2-pro", "flux-2-pro"), + ], +) +def test_flux2_uses_model_specific_provider_url(model: str, provider_path: str): + url = AzureChatCompletion().create_azure_base_url( + azure_client_params={ + "azure_endpoint": "https://example.services.ai.azure.com/", + "api_version": "preview", + }, + model=model, + ) + + assert ( + url == f"https://example.services.ai.azure.com/providers/blackforestlabs/v1/{provider_path}?api-version=preview" + ) + + +def test_flux2_flex_maps_openai_and_provider_parameters(): + config = AzureFoundryFluxImageGenerationConfig() + mapped_params = config.map_openai_params( + non_default_params={ + "n": 2, + "size": "1536x1024", + "guidance": 4.5, + "steps": 32, + "output_format": "jpeg", + }, + optional_params={}, + model="FLUX.2-flex", + drop_params=False, + ) + url = config.get_flux2_image_generation_url( + api_base="https://example.services.ai.azure.com", + model="FLUX.2-flex", + api_version="preview", + ) + request = azure_deployment_image_generation_json_body( + api_base=url, + data={"model": "FLUX.2-flex", "prompt": "A red fox", **mapped_params}, + deployment_name="FLUX.2-flex", + ) + + assert request == { + "model": "FLUX.2-flex", + "prompt": "A red fox", + "num_images": 2, + "width": 1536, + "height": 1024, + "guidance": 4.5, + "steps": 32, + "output_format": "jpeg", + } + + +def test_flux2_flex_rejects_invalid_size(): + with pytest.raises(ValueError, match="Expected 'WxH'"): + AzureFoundryFluxImageGenerationConfig().map_openai_params( + non_default_params={"size": "large"}, + optional_params={}, + model="FLUX.2-flex", + drop_params=False, + ) + + +def test_flux2_flex_model_info(): + model_info = litellm.get_model_info( + model="FLUX.2-flex", + custom_llm_provider="azure_ai", + ) + catalog_info = litellm.model_cost["azure_ai/FLUX.2-flex"] + + assert model_info["mode"] == "image_generation" + assert model_info["max_input_tokens"] == 32000 + assert model_info["max_tokens"] == 32000 + assert model_info["supported_endpoints"] == ["/v1/images/generations", "/v1/images/edits"] + assert catalog_info["input_cost_per_pixel"] == 5e-08 + assert catalog_info["supported_modalities"] == ["text", "image"] + assert catalog_info["supported_output_modalities"] == ["image"] + + +def test_flux2_flex_cost_uses_generated_megapixels(): + response = ImageResponse( + data=[ + ImageObject(url="https://example.com/one.png"), + ImageObject(url="https://example.com/two.png"), + ] + ) + + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="FLUX.2-flex", + completion_response=response, + custom_llm_provider="azure_ai", + size="2048x1024", + call_type="image_generation", + ) + + assert cost == pytest.approx(5e-08 * 2048 * 1024 * 2) + + +@pytest.mark.parametrize("model", ("FLUX-1.1-pro", "FLUX.1-Kontext-pro")) +def test_flux1_preserves_existing_openai_parameters(model: str): + params: Final = {"n": 2, "size": "1536x1024", "quality": "high", "user": "test-user"} + + mapped: Final = AzureFoundryFluxImageGenerationConfig().map_openai_params( + non_default_params=params, + optional_params={}, + model=model, + drop_params=False, + ) + + assert mapped == params + + +@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024})) +def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensions: Mapping[str, int | str]): + params: Final = AzureFoundryFluxImageGenerationConfig().map_openai_params( + non_default_params={"n": 2, **dimensions}, + optional_params={}, + model="FLUX.2-flex", + drop_params=False, + ) + response: Final = get_azure_image_generation_config("FLUX.2-flex").transform_image_generation_response( + model="FLUX.2-flex", + raw_response=httpx.Response(200, json={"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]}), + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={"prompt": "A red fox", **params}, + optional_params=params, + litellm_params={}, + encoding=None, + ) + + assert litellm.completion_cost( + model="azure_ai/FLUX.2-flex", + completion_response=response, + optional_params=params, + call_type="image_generation", + ) == pytest.approx(5e-08 * 2048 * 1024 * 2) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b26f5e25b6f..6b7337ea8d1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35343,7 +35343,7 @@ export interface components { default_model?: string | null; /** * Deployment Affinity - * @description When True and a session_id is resolvable on the request, pin the deployment chosen inside each routed model group and reuse it whenever the session returns to that group, without pinning which group the session routes to. Independent of session_affinity, which pins the model group instead (and always carries this deployment pin with it): with session_affinity off, every turn is still classified on its own merits while a session that escalates to a stronger tier and comes back still lands on the deployment it used before, which is what keeps a provider prompt cache warm. Pins are held per model group, so switching tiers does not disturb the pin left behind in the previous group. On by default because re-shuffling a conversation across deployments of the same model discards that cache for no benefit; set False to keep every turn load-balanced across the group, which is what a deployment set with tight per-deployment rate limits wants. Inert when no session_id is resolvable, since there is nothing to key a pin on, and suppressed when plugins are configured, for the same reason session_affinity is. + * @description When True and a client session_id is resolvable, reuse the session's chosen model for each classified tier and its deployment within each model group. With session_affinity off, every turn is still classified: moving to another tier leaves the previous tier's model pin intact for a later return. Pins yield to current candidate, context, modality, and availability constraints. Adaptive selection chooses the initial model from its eligible pool, then reuses that choice per tier. This reduces avoidable provider prompt-cache misses; it does not guarantee cache hits. Set False to select models and load-balance deployments on every turn, unless session_affinity or user_turn classification requires a pin. Inert without a client session_id and suppressed when plugins are configured. * @default true */ deployment_affinity: boolean; @@ -35487,7 +35487,7 @@ export interface components { session_affinity: boolean; /** * Session Affinity Ttl Seconds - * @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity deployment pin, so it measures idle time for the session's routing decisions rather than total session length + * @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity per-tier model and deployment pins, so it measures idle time for the session's routing decisions rather than total session length * @default 3600 */ session_affinity_ttl_seconds: number; From ff1e2a02ba9b7bafee87896edcbf935749bf59c3 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:33:58 -0500 Subject: [PATCH 015/102] fix(proxy): preserve CI-compatible OpenAPI snapshot formatting --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 81889e4da0c..fb2f014e3d8 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -18974,7 +18974,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 3821e5ace617f43122e757ec0d922c6f4a8c9b6a Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:43:24 -0500 Subject: [PATCH 016/102] fix(azure_ai): specify FLUX parameter mapping return type --- litellm/llms/azure_ai/image_generation/flux_transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index b10e8a6f35e..b9d5e11ff2a 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -102,7 +102,7 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): optional_params: Mapping[str, object], model: str, drop_params: bool, - ) -> dict: # mutable-ok: inherited config contract returns a dict + ) -> dict[str, object]: # mutable-ok: inherited config contract returns a dict if not self.is_flux2_model(model): return super().map_openai_params( non_default_params=dict(non_default_params), From 4595b4f62f209d3d71734e8f1f2692f9a726e9a1 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 13:01:12 -0500 Subject: [PATCH 017/102] fix(azure-ai): coerce FLUX controls and preserve response dimensions --- .../image_generation/flux_transformation.py | 5 +++++ .../image_generation/gpt_transformation.py | 9 ++++++++- .../test_azure_ai_image_edit_transformation.py | 6 +++--- .../test_azure_ai_flux2_image_generation.py | 18 ++++++++++++++++++ 4 files changed, 34 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index b9d5e11ff2a..997205b1fc3 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -85,6 +85,11 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): @staticmethod def _map_parameter(name: str, value: object) -> tuple[tuple[str, object], ...]: + if isinstance(value, str): + if name in ("n", "num_images", "width", "height", "steps", "seed", "safety_tolerance"): + return (("num_images" if name == "n" else name, int(value)),) + if name == "guidance": + return ((name, float(value)),) if name == "n": return (("num_images", value),) if name != "size": diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 090b2eba387..8dc4d8953ea 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -82,7 +82,14 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): ) # set optional params - image_response.size = image_response.size or optional_params.get("size", "1024x1024") + width: Final = optional_params.get("width") + height: Final = optional_params.get("height") + requested_size: Final = ( + f"{width}x{height}" + if isinstance(width, int) and isinstance(height, int) + else optional_params.get("size", "1024x1024") + ) + image_response.size = image_response.size or requested_size image_response.quality = image_response.quality or optional_params.get("quality", "high") image_response.output_format = image_response.output_format or optional_params.get("output_format", "png") diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 94b23c5f1c2..d74afa88a6b 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -144,7 +144,7 @@ def test_flux2_image_edit_rejects_too_many_references(model: str, reference_imag ) -@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024})) +@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024}, {"width": "2048", "height": "1024"})) @pytest.mark.usefixtures("local_model_cost_map") def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[str, int | str]): def respond(request: httpx.Request) -> httpx.Response: @@ -170,8 +170,8 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ api_base="https://example.services.ai.azure.com", client=client, n=2, - guidance=4.5, - steps=32, + guidance="4.5", + steps="32", **dimensions, ) diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index 07026cf4309..c1d46b7e919 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -170,3 +170,21 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi optional_params=params, call_type="image_generation", ) == pytest.approx(5e-08 * 2048 * 1024 * 2) + + +def test_flux2_response_preserves_mapped_dimensions(): + config = AzureFoundryFluxImageGenerationConfig() + params = config.map_openai_params( + non_default_params={"size": "2048x1024"}, optional_params={}, model="FLUX.2-flex", drop_params=False + ) + response = config.transform_image_generation_response( + model="FLUX.2-flex", + raw_response=httpx.Response(200, json={"data": [{"b64_json": "aW1n"}]}), + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={"prompt": "A landscape"}, + optional_params=params, + litellm_params={}, + encoding=None, + ) + assert response.size == "2048x1024" From c8a2d8c3496ba643aa930b16982382d9a5d76d8c Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:00:36 +0000 Subject: [PATCH 018/102] feat(proxy): add LiteLLM_DailyGlobalSpend key-free rollup for the usage dashboard Adds a daily spend table without api_key or user_id, written atomically alongside LiteLLM_DailyUserSpend from the batched writer, reconciled from history by a scheduled job that advances a marker in LiteLLM_Config, and read by the key-free arm of the aggregated usage query once the marker covers the requested range. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 33 ++ .../litellm_proxy_extras/schema.prisma | 29 ++ litellm/constants.py | 3 + litellm/proxy/db/daily_spend_bulk_upsert.py | 98 +++-- litellm/proxy/db/db_spend_update_writer.py | 7 +- .../common_daily_activity.py | 38 +- litellm/proxy/proxy_server.py | 43 ++ litellm/proxy/schema.prisma | 29 ++ .../daily_global_spend_rollup.py | 235 +++++++++++ schema.prisma | 29 ++ .../proxy/db/test_daily_spend_bulk_upsert.py | 150 +++++++ .../proxy/db/test_db_spend_update_writer.py | 101 ++++- .../test_common_daily_activity.py | 154 ++++++- .../proxy/proxy_server/test_lifecycle.py | 48 +++ .../test_daily_global_spend_rollup.py | 382 ++++++++++++++++++ 15 files changed, 1347 insertions(+), 32 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql create mode 100644 litellm/proxy/spend_tracking/daily_global_spend_rollup.py create mode 100644 tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql new file mode 100644 index 00000000000..1d6cdea0c7b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql @@ -0,0 +1,33 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGlobalSpend" ( + "id" TEXT NOT NULL, + "date" TEXT NOT NULL, + "model" TEXT, + "model_group" TEXT, + "custom_llm_provider" TEXT, + "mcp_namespaced_tool_name" TEXT, + "endpoint" TEXT, + "prompt_tokens" BIGINT NOT NULL DEFAULT 0, + "completion_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0, + "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0, + "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "api_requests" BIGINT NOT NULL DEFAULT 0, + "successful_requests" BIGINT NOT NULL DEFAULT 0, + "failed_requests" BIGINT NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyGlobalSpend_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGlobalSpend_date_idx" ON "LiteLLM_DailyGlobalSpend"("date"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyGlobalSpend_date_model_model_group_custom_llm__key" ON "LiteLLM_DailyGlobalSpend"("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8072df5aa5b..5d433e916d6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -809,6 +809,35 @@ model LiteLLM_DailyUserSpend { @@index([endpoint]) } +// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view +model LiteLLM_DailyGlobalSpend { + id String @id @default(uuid()) + date String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@index([date]) +} + // Track daily organization spend metrics per model and key model LiteLLM_DailyOrganizationSpend { id String @id @default(uuid()) diff --git a/litellm/constants.py b/litellm/constants.py index 565c6433c6e..1bf3a150aeb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2034,6 +2034,9 @@ PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90 # Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide # expiry cannot produce an alert too large for the channel delivering it. PTU_LAPSED_ALERT_LIMIT: Final[int] = 10 +DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID: Final[str] = "daily_global_spend_reconcile_job" +DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS: Final[int] = 3600 +DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM: Final[str] = "daily_global_spend_reconciled_through" # Slack allowed when deciding a sentinel row is stale. The row's updated_at and the # run's cutoff are stamped by different hosts, so clock skew between them must not let # one run delete a charge another just wrote. A stale row is hours old and a concurrent diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index a143643577e..c83043101eb 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -25,29 +25,41 @@ SpendRow = Mapping[str, object] @dataclass(frozen=True, slots=True) class DailySpendTable: - """The physical table behind one entity's daily rollup.""" + """A daily rollup table and the unique constraint its upserts arbitrate on.""" name: str - entity_id_column: str + key_columns: tuple[str, ...] carries_request_id: bool = False -DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( - { - "user": DailySpendTable(name="LiteLLM_DailyUserSpend", entity_id_column="user_id"), - "team": DailySpendTable(name="LiteLLM_DailyTeamSpend", entity_id_column="team_id"), - "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), - "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), - "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), - "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), - } -) - # The unique constraint's columns after the entity id, in constraint order. A NULL can # never match itself in a unique index, so every one of these is normalized to '': the # conflict target has to be NULL-free or the row is re-inserted on every single flush. _KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") + +def _entity_table(name: str, entity_id_column: str, carries_request_id: bool = False) -> DailySpendTable: + return DailySpendTable( + name=name, key_columns=(entity_id_column, *_KEY_COLUMNS), carries_request_id=carries_request_id + ) + + +DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( + { + "user": _entity_table("LiteLLM_DailyUserSpend", "user_id"), + "team": _entity_table("LiteLLM_DailyTeamSpend", "team_id"), + "org": _entity_table("LiteLLM_DailyOrganizationSpend", "organization_id"), + "end_user": _entity_table("LiteLLM_DailyEndUserSpend", "end_user_id"), + "agent": _entity_table("LiteLLM_DailyAgentSpend", "agent_id"), + "tag": _entity_table("LiteLLM_DailyTagSpend", "tag", carries_request_id=True), + } +) + +GLOBAL_SPEND_TABLE: Final = DailySpendTable( + name="LiteLLM_DailyGlobalSpend", + key_columns=("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"), +) + _COUNTER_COLUMNS: Final = ( "prompt_tokens", "completion_tokens", @@ -92,7 +104,7 @@ def _as_float(value: object) -> float: def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" - return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) + return tuple(_as_text(transaction.get(column)) for column in table.key_columns) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -130,7 +142,11 @@ def _row_params( return ( str(uuid.uuid4()), *key, - None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")), + *( + () + if "model_group" in table.key_columns + else (None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")),) + ), *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), @@ -140,26 +156,25 @@ def _row_params( def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - table.entity_id_column, - *_KEY_COLUMNS, - "model_group", + *table.key_columns, + *(() if "model_group" in table.key_columns else ("model_group",)), *_COUNTER_COLUMNS, *_SPEND_COLUMNS, *(("request_id",) if table.carries_request_id else ()), ) -def build_bulk_upsert( +def _upsert_statement( table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]], -) -> tuple[str, tuple[SqlValue, ...]]: - """The single statement writing one merged batch, plus its positional arguments.""" + first_param: int, +) -> str: columns: Final = _insert_columns(table) quoted_table: Final = f'"{table.name}"' rows: Final = ", ".join( "(" + ", ".join( - f"${row_index * len(columns) + offset + 1}::{_CASTS.get(column, 'text')}" + f"${first_param + row_index * len(columns) + offset}::{_CASTS.get(column, 'text')}" for offset, column in enumerate(columns) ) + ", (NOW() AT TIME ZONE 'UTC'))" @@ -176,11 +191,44 @@ def build_bulk_upsert( if table.carries_request_id else "" ) - sql: Final = ( + return ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted(table.key_columns)}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) - return sql, tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) + + +def _params(table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]]) -> tuple[SqlValue, ...]: + return tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) + + +def build_bulk_upsert( + table: DailySpendTable, + batch: Sequence[tuple[tuple[str, ...], SpendRow]], +) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one merged batch, plus its positional arguments.""" + return _upsert_statement(table, batch, first_param=1), _params(table, batch) + + +def build_bulk_upsert_with_global_rollup( + table: DailySpendTable, + batch: Sequence[tuple[tuple[str, ...], SpendRow]], +) -> tuple[str, tuple[SqlValue, ...]]: + """One statement writing a batch to its table and, atomically, its key-free rollup + to ``LiteLLM_DailyGlobalSpend``. + + A data-modifying CTE runs both inserts in the same snapshot and transaction, so a + batch that lands in one table lands in both and a retried deadlock replays both. + Postgres does not order the CTE against the main statement, so two writers can still + deadlock across the tables; the caller's deadlock retry covers that, and each insert + takes its own rows in key order so same-table lock order stays deterministic. + """ + global_batch: Final = merge_by_conflict_key(GLOBAL_SPEND_TABLE, tuple(row for _, row in batch)) + entity_params: Final = _params(table, batch) + sql: Final = ( + f"WITH entity_rows AS (\n{_upsert_statement(table, batch, first_param=1)}\nRETURNING 1)\n" + f"{_upsert_statement(GLOBAL_SPEND_TABLE, global_batch, first_param=len(entity_params) + 1)}" + ) + return sql, (*entity_params, *_params(GLOBAL_SPEND_TABLE, global_batch)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index eaa03c5d7f7..d5c839a9be8 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -46,6 +46,7 @@ from litellm.proxy._types import ( from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, + build_bulk_upsert_with_global_rollup, merge_by_conflict_key, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -1939,7 +1940,11 @@ class DBSpendUpdateWriter: merged_batch = merge_by_conflict_key( table=table, transactions=tuple(transactions_to_process.values()) ) - sql, params = build_bulk_upsert(table=table, batch=merged_batch) + sql, params = ( + build_bulk_upsert_with_global_rollup(table=table, batch=merged_batch) + if entity_type == "user" + else build_bulk_upsert(table=table, batch=merged_batch) + ) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: # Log detailed error information for debugging batch upsert failures diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 8a3ba196ab2..f1d78dca201 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -11,6 +11,8 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE +from litellm.proxy.spend_tracking.daily_global_spend_rollup import reconciled_through from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, recover_double_hashed_key_metadata, @@ -734,6 +736,30 @@ def _rollup_metric_select(table_name: str) -> str: _MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" +async def key_free_source_table(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: + """The table the key-free arm reads from, when the global rollup can answer instead of the per-key table. + + Only an unfiltered read of the user table has the same rows as ``LiteLLM_DailyGlobalSpend``, + and only through the day the reconcile marker has reached: the writer keeps that day + current, later days are covered once the next run advances the marker. + """ + if query["table_name"] != "litellm_dailyuserspend": + return None + if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]: + return None + _, adjusted_end = _adjust_dates_for_timezone( + query["start_date"], query["end_date"], query["timezone_offset_minutes"], query["include_current_utc_day"] + ) + try: + marker: Final = await reconciled_through(prisma_client) + except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read + verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc) + return None + if marker is None or adjusted_end > marker: + return None + return GLOBAL_SPEND_TABLE.name + + def _build_aggregated_sql_query( *, table_name: str, @@ -746,13 +772,16 @@ def _build_aggregated_sql_query( exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path timezone_offset_minutes: int | None = None, include_current_utc_day: bool = False, + key_free_table: str | None = None, ) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params """Build the GROUPING SETS query for aggregated daily activity. One statement, two UNION ALL arms over the same WHERE clause. The first arm is key-free: grand total, per-date totals and the (date, model / model_group / provider / mcp / endpoint) rollups, so its row count never grows with the number - of keys. The second arm emits the (date, , api_key) rollups for the + of keys; it reads ``key_free_table`` when given (the global rollup, whose row count + never grew with the number of keys to begin with) and the entity table otherwise. + The second arm emits the (date, , api_key) rollups for the USAGE_TOP_API_KEYS_LIMIT highest-spend keys only. Both arms share the 7-bit group_level bitmask (date, api_key, model, model_group, provider, mcp, endpoint). @@ -778,6 +807,7 @@ def _build_aggregated_sql_query( ) sentinel_param: Final = f"${len(where_params) + 1}" metric_select: Final = _rollup_metric_select(table_name) + key_free_source: Final = key_free_table or pg_table # TODO: drop the successful_requests/failed_requests aggregates (and the # total_successful_requests metadata they feed) once the admin UI reads SGR @@ -796,7 +826,7 @@ def _build_aggregated_sql_query( | GROUPING(model, {_MODEL_GROUP_EXPR}, custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level,{metric_select} - FROM "{pg_table}" + FROM "{key_free_source}" WHERE {where_clause} GROUP BY GROUPING SETS ( (date), @@ -1387,7 +1417,9 @@ async def get_daily_activity_aggregated( timezone_offset_minutes=timezone_offset_minutes, include_current_utc_day=include_current_utc_day, ) - sql_query, sql_params = _build_aggregated_sql_query(**query_kwargs) + sql_query, sql_params = _build_aggregated_sql_query( + **query_kwargs, key_free_table=await key_free_source_table(prisma_client, query_kwargs) + ) entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None raw_rows, raw_entity_rows = await asyncio.gather( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f63e088ebf7..ed8f6886734 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -259,6 +259,7 @@ from litellm.constants import ( APSCHEDULER_MISFIRE_GRACE_TIME, APSCHEDULER_REPLACE_EXISTING, CLI_SSO_SESSION_TTL_SECONDS, + DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, @@ -662,6 +663,9 @@ from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( + run_scheduled_daily_global_spend_reconcile, +) from litellm.proxy.spend_tracking.spend_counter_batch import ( PendingSpendIncrement, active_spend_counter_batch, @@ -9970,6 +9974,12 @@ class ProxyStartupEvent: await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) + cls._initialize_daily_global_spend_reconcile_job( + scheduler=scheduler, + proxy_logging_obj=proxy_logging_obj, + prisma_client=prisma_client, + ) + ### PTU DAILY ROLLUP ### from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, @@ -10311,6 +10321,39 @@ class ProxyStartupEvent: "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)" ) + @classmethod + def _initialize_daily_global_spend_reconcile_job( + cls, + scheduler: AsyncIOScheduler, + proxy_logging_obj: ProxyLogging, + prisma_client: PrismaClient, + ) -> None: + async def alert(message: str) -> None: + await proxy_logging_obj.alerting_handler( + message=message, + level="High", + alert_type=AlertType.failed_tracking_spend, + ) + + async def reconcile() -> None: + await run_scheduled_daily_global_spend_reconcile( + prisma_client, + pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, + alert=alert, + ) + + scheduler.add_job( + reconcile, + "cron", + hour=0, + minute=30, + timezone="UTC", + id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + next_run_time=datetime.now(timezone.utc) + timedelta(minutes=2), + ) + @classmethod async def _initialize_slack_alerting_jobs( cls, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8072df5aa5b..5d433e916d6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -809,6 +809,35 @@ model LiteLLM_DailyUserSpend { @@index([endpoint]) } +// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view +model LiteLLM_DailyGlobalSpend { + id String @id @default(uuid()) + date String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@index([date]) +} + // Track daily organization spend metrics per model and key model LiteLLM_DailyOrganizationSpend { id String @id @default(uuid()) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py new file mode 100644 index 00000000000..9d344421332 --- /dev/null +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -0,0 +1,235 @@ +"""Reconcile ``LiteLLM_DailyGlobalSpend`` from ``LiteLLM_DailyUserSpend``, one day per transaction. + +The spend writer keeps both tables in step from the moment it is deployed; this job rolls up +the days before that and records how far it has reached in ``LiteLLM_Config`` so usage reads +know when the global table can answer for a date range. It runs as a background cron, never +in a Prisma migration, since on a large deployment the aggregate is minutes of work. +""" + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import date, datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, + DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, +) +from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE +from litellm.repositories.config_repository import ConfigRepository + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + from litellm.proxy.utils import PrismaClient + +_DAY_TRANSACTION_TIMEOUT: Final = timedelta(minutes=10) +_REPLAY_DAYS: Final = 1 +_METRIC_COLUMNS: Final = ( + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "compression_saved_tokens", + "api_requests", + "successful_requests", + "failed_requests", + "compression_savings_spend", + "prompt_caching_savings_spend", + "gateway_injected_caching_savings_spend", + "autorouter_savings_spend", + "spend", +) + + +def _quoted(columns: tuple[str, ...]) -> str: + return ", ".join(f'"{column}"' for column in columns) + + +def _reconcile_day_sql() -> str: + key_columns: Final = GLOBAL_SPEND_TABLE.key_columns + normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in key_columns) + sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS) + overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS) + return ( + f'INSERT INTO "{GLOBAL_SPEND_TABLE.name}" ("id", {_quoted(key_columns)}, {_quoted(_METRIC_COLUMNS)}, ' + '"updated_at")\n' + f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n" + 'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n' + f"GROUP BY {normalized_keys}\n" + f"ON CONFLICT ({_quoted(key_columns)}) DO UPDATE SET {overwrite}, " + "\"updated_at\" = (NOW() AT TIME ZONE 'UTC')" + ) + + +RECONCILE_DAY_SQL: Final = _reconcile_day_sql() +_LOCK_GLOBAL_TABLE_SQL: Final = f'LOCK TABLE "{GLOBAL_SPEND_TABLE.name}" IN EXCLUSIVE MODE' +_PENDING_DAYS_SQL: Final = ( + 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" >= $1 AND "date" <= $2 ORDER BY "date"' +) + + +class ReconciledThrough(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + reconciled_through: str + + +class _MarkerRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + + param_value: object = None + + +class _DateRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + date: str + + +@dataclass(frozen=True, slots=True) +class ReconcileResult: + days_reconciled: tuple[str, ...] + reconciled_through: str | None + failed_day: str | None = None + + +def _marker_from_param_value(value: object) -> str | None: + try: + parsed: Final = ( + ReconciledThrough.model_validate_json(value) + if isinstance(value, str) + else ReconciledThrough.model_validate(value) + ) + except ValidationError: + return None + return parsed.reconciled_through + + +async def reconciled_through(prisma_client: "PrismaClient") -> str | None: + """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" + from litellm.proxy.utils import get_config_param + + row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value) + + +async def _record_reconciled_through(prisma_client: "PrismaClient", day: str) -> None: + from litellm.proxy.utils import invalidate_config_param + + await ConfigRepository(prisma_client).set_param( + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ReconciledThrough(reconciled_through=day).model_dump_json() + ) + await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + +def _first_pending_day(marker: str | None) -> str: + if marker is None: + return "" + return (date.fromisoformat(marker) - timedelta(days=_REPLAY_DAYS)).isoformat() + + +async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: + """Every UTC day through today still to roll up, oldest first; the marker day and the one + before it are replayed so rows flushed by a pre-writer pod during a rolling deploy are folded in.""" + marker: Final = await reconciled_through(prisma_client) + rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), today.isoformat()) + return tuple(sorted({*(_DateRow.model_validate(row).date for row in rows), today.isoformat()})) + + +async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: + """Rewrite one day of the global table from the per-key sums; the table lock keeps the + writer's increments out between the aggregate and the overwrite so none are lost.""" + async with prisma_client.db.tx(timeout=_DAY_TRANSACTION_TIMEOUT) as transaction: + await transaction.execute_raw(_LOCK_GLOBAL_TABLE_SQL) + await transaction.execute_raw(RECONCILE_DAY_SQL, day) + + +async def run_daily_global_spend_reconcile( + prisma_client: "PrismaClient", + today: date | None = None, +) -> ReconcileResult: + """Roll up every pending day, advancing the marker after each; a failing day stops the run + with the marker on the last good day so the next run resumes there.""" + effective_today: Final = today or datetime.now(timezone.utc).date() + days: Final = await pending_days(prisma_client, effective_today) + done: Final = await _reconcile_until_failure(prisma_client, days) + failed: Final = days[len(done)] if len(done) < len(days) else None + marker: Final = done[-1] if done else await reconciled_through(prisma_client) + return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=failed) + + +async def _reconcile_until_failure(prisma_client: "PrismaClient", days: tuple[str, ...]) -> tuple[str, ...]: + for index, day in enumerate(days): + if not await _reconcile_and_record(prisma_client, day): + return days[:index] + return days + + +async def _reconcile_and_record(prisma_client: "PrismaClient", day: str) -> bool: + try: + await reconcile_day(prisma_client, day) + await _record_reconciled_through(prisma_client, day) + except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done + verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) + return False + return True + + +async def run_scheduled_daily_global_spend_reconcile( + prisma_client: "PrismaClient", + pod_lock_manager: "PodLockManager | None" = None, + alert: Callable[[str], Awaitable[None]] | None = None, + today: date | None = None, +) -> ReconcileResult | None: + """Run the reconcile under a cross-pod lock so one proxy does the work; the lock only saves + effort (each day is an idempotent rewrite), so an unreachable Redis runs unguarded rather than skipping.""" + redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache + if pod_lock_manager is None or redis_cache is None: + return await _run_and_alert(prisma_client, alert=alert, today=today) + + acquired: Final = await pod_lock_manager.acquire_lock( + cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, ttl=DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS + ) + if not acquired and await _lock_is_held(pod_lock_manager, redis_cache): + verbose_proxy_logger.info("Daily global spend reconcile: another pod holds the lock, skipping this run") + return None + try: + return await _run_and_alert(prisma_client, alert=alert, today=today) + finally: + if acquired: + await pod_lock_manager.release_lock(cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) + + +async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool: + try: + lock_key: Final = pod_lock_manager.get_redis_lock_key(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) + return bool(await redis_cache.async_get_cache(lock_key)) + except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the run + verbose_proxy_logger.warning("Daily global spend reconcile: could not read the lock: %s", exc) + return False + + +async def _run_and_alert( + prisma_client: "PrismaClient", + *, + alert: Callable[[str], Awaitable[None]] | None, + today: date | None, +) -> ReconcileResult: + result: Final = await run_daily_global_spend_reconcile(prisma_client, today=today) + if result.days_reconciled: + verbose_proxy_logger.info( + "Daily global spend reconcile: rolled up %d day(s), reconciled through %s", + len(result.days_reconciled), + result.reconciled_through, + ) + if result.failed_day is not None and alert is not None: + await alert( + f"Daily global spend reconcile stopped at {result.failed_day}; usage totals keep reading the per-key " + f"table for ranges past {result.reconciled_through or 'the beginning'} until the next run succeeds." + ) + return result diff --git a/schema.prisma b/schema.prisma index 8072df5aa5b..5d433e916d6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -809,6 +809,35 @@ model LiteLLM_DailyUserSpend { @@index([endpoint]) } +// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view +model LiteLLM_DailyGlobalSpend { + id String @id @default(uuid()) + date String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@index([date]) +} + // Track daily organization spend metrics per model and key model LiteLLM_DailyOrganizationSpend { id String @id @default(uuid()) diff --git a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py index c1efb3e7220..cc443a2cfe5 100644 --- a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py @@ -1,12 +1,19 @@ """Tests for the single-statement daily spend upsert (LIT-5291).""" +import pathlib import re +from typing import Final +import psycopg import pytest +from psycopg.rows import dict_row +from pytest_postgresql import factories from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, + GLOBAL_SPEND_TABLE, build_bulk_upsert, + build_bulk_upsert_with_global_rollup, conflict_key, merge_by_conflict_key, ) @@ -185,3 +192,146 @@ async def test_writer_survives_a_transaction_whose_key_columns_are_null(): _, params = prisma_client.db.statements[0] assert None not in params[:9] assert transactions == {} + + +def user_txn(**overrides): + txn = {**tag_txn(), "user_id": "u-1", **overrides} + del txn["tag"] + del txn["request_id"] + return txn + + +def _bound_rows(insert_sql: str, params: tuple[object, ...]) -> list[dict[str, object]]: + """Each VALUES row of one INSERT as a column -> bound value mapping, consuming params in order.""" + header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", insert_sql) + assert header is not None, insert_sql + columns = [c.strip('"') for c in header.group(1).split(", ") if c != '"updated_at"'] + row_count = insert_sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") + return [dict(zip(columns, params[i * len(columns) : (i + 1) * len(columns)])) for i in range(row_count)] + + +def test_global_rollup_folds_every_key_and_user_into_one_row_per_dimension_tuple(): + """The global table has no api_key or user_id, so a batch spread over many keys and + users must collapse to one row per (date, model, group, provider, mcp, endpoint).""" + batch = merge_by_conflict_key( + USER_TABLE, + tuple(user_txn(user_id=f"u-{i}", api_key=f"sk-{i}", spend=1.0, api_requests=1) for i in range(5)) + + (user_txn(user_id="u-0", api_key="sk-0", model="claude", spend=10.0, api_requests=3),), + ) + + sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) + + entity_insert, global_insert = sql.split("RETURNING 1)") + entity_rows = _bound_rows(entity_insert, params) + global_rows = _bound_rows(global_insert, params[len(entity_rows) * len(entity_rows[0]) :]) + assert len(entity_rows) == 6 + assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in global_insert + assert [(r["model"], r["spend"], r["api_requests"]) for r in global_rows] == [ + ("claude", 10.0, 3), + ("gpt-4o-mini", 5.0, 5), + ] + assert all("api_key" not in r and "user_id" not in r for r in global_rows) + conflict = re.search(r"ON CONFLICT \(([^)]*)\)", global_insert) + assert conflict is not None + assert conflict.group(1) == ", ".join(f'"{c}"' for c in GLOBAL_SPEND_TABLE.key_columns) + + +def test_global_rollup_params_follow_the_entity_params_in_one_placeholder_sequence(): + """Both inserts bind from one flat tuple, so the global arm's placeholders must start + exactly where the entity arm's stop or every value lands one column off.""" + batch = merge_by_conflict_key(USER_TABLE, (user_txn(),)) + + sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) + + placeholders = [int(n) for n in re.findall(r"\$(\d+)::", sql)] + assert placeholders == list(range(1, len(params) + 1)) + + +_bulk_upsert_postgresql_proc: Final = factories.postgresql_proc() +_bulk_upsert_postgresql: Final = factories.postgresql("_bulk_upsert_postgresql_proc") + +_MIGRATIONS_DIR: Final = ( + pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" +) +_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" + +_DAILY_USER_SPEND_DDL: Final = """ + CREATE TABLE "LiteLLM_DailyUserSpend" ( + id TEXT PRIMARY KEY, + user_id TEXT, + date TEXT NOT NULL, + api_key TEXT NOT NULL, + model TEXT, + model_group TEXT, + custom_llm_provider TEXT, + mcp_namespaced_tool_name TEXT, + endpoint TEXT, + prompt_tokens BIGINT DEFAULT 0, + completion_tokens BIGINT DEFAULT 0, + cache_read_input_tokens BIGINT DEFAULT 0, + cache_creation_input_tokens BIGINT DEFAULT 0, + compression_saved_tokens BIGINT DEFAULT 0, + compression_savings_spend DOUBLE PRECISION DEFAULT 0, + prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, + spend DOUBLE PRECISION DEFAULT 0, + api_requests BIGINT DEFAULT 0, + successful_requests BIGINT DEFAULT 0, + failed_requests BIGINT DEFAULT 0, + created_at TIMESTAMP DEFAULT now(), + updated_at TIMESTAMP, + UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) + ) +""" + + +def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: + converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) + conn.execute( + converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query + {f"p{i}": v for i, v in enumerate(params, start=1)}, + ) + conn.commit() + + +def test_global_rollup_equals_the_per_key_sums_after_repeated_flushes(_bulk_upsert_postgresql: psycopg.Connection): + """Against real Postgres and the shipped migration: two flushes of a mixed batch leave + the global table exactly equal to the per-key table summed over user and key, with the + NULL and '' spellings of a dimension folded into one row.""" + conn: Final = _bulk_upsert_postgresql + conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal + conn.commit() + + batch = merge_by_conflict_key( + USER_TABLE, + ( + user_txn(user_id="u-1", api_key="sk-1", spend=1.0, prompt_tokens=10), + user_txn(user_id="u-2", api_key="sk-2", spend=2.0, prompt_tokens=20), + user_txn(user_id="u-1", api_key="sk-3", model=None, custom_llm_provider=None, spend=4.0), + user_txn(user_id="u-3", api_key="sk-4", model="", custom_llm_provider="", spend=8.0), + ), + ) + sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) + _execute_dollar_sql(conn, sql, params) + _execute_dollar_sql(conn, sql, params) + + with conn.cursor(row_factory=dict_row) as cur: + global_rows = cur.execute( + 'SELECT model, spend, prompt_tokens, api_requests FROM "LiteLLM_DailyGlobalSpend" ORDER BY model' + ).fetchall() + per_key = cur.execute( + """ + SELECT COALESCE(model, '') AS model, SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, + SUM(api_requests) AS api_requests + FROM "LiteLLM_DailyUserSpend" GROUP BY COALESCE(model, '') ORDER BY 1 + """ + ).fetchall() + + assert [row["model"] for row in global_rows] == ["", "gpt-4o-mini"] + assert [(r["model"], r["spend"], int(r["prompt_tokens"]), int(r["api_requests"])) for r in global_rows] == [ + (r["model"], float(r["spend"]), int(r["prompt_tokens"]), int(r["api_requests"])) for r in per_key + ] + assert global_rows[0]["spend"] == pytest.approx(24.0) + assert global_rows[1]["spend"] == pytest.approx(6.0) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 5e977712a1e..d8a9013398e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -254,14 +254,19 @@ class _RecordingPrisma: def _row_values(statement: Statement, column: str) -> list[object]: - """Every row's value for one column, read out of the flat parameter tuple.""" + """Every row's value for one column of the first INSERT, read out of the flat parameter tuple. + + The user-table statement chains a global rollup INSERT after its own, so the row count + comes from the first INSERT's VALUES rather than from the parameter count. + """ sql, params = statement header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", sql) assert header is not None, sql columns = header.group(1).split(", ") stride = len(columns) - 1 # updated_at is inlined, not bound offset = columns.index(f'"{column}"') - return [params[row * stride + offset] for row in range(len(params) // stride)] + rows = sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") + return [params[row * stride + offset] for row in range(rows)] @pytest.mark.asyncio @@ -1463,6 +1468,98 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry(): assert daily_spend_transactions == expected +def _entity_txn(entity_field: str, entity_id: str, api_key: str) -> dict[str, object]: + txn = _daily_txn() + del txn["user_id"] + return {**txn, entity_field: entity_id, "api_key": api_key} + + +@pytest.mark.asyncio +async def test_user_flush_writes_the_global_rollup_in_the_same_statement(): + """The user flush is the one place per-key spend becomes key-free spend, so a batch spread + over many keys must land in LiteLLM_DailyGlobalSpend as one row in the same statement. + A separate statement would let a crash between the two leave the tables out of sync.""" + prisma_client = _RecordingPrisma() + txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(4)} + + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=MagicMock(), + daily_spend_transactions=txns, + entity_type="user", + entity_id_field="user_id", + ) + + assert len(prisma_client.db.statements) == 1 + sql, params = prisma_client.db.statements[0] + assert sql.count('INSERT INTO "LiteLLM_DailyUserSpend"') == 1 + assert sql.count('INSERT INTO "LiteLLM_DailyGlobalSpend"') == 1 + assert sql.index('"LiteLLM_DailyUserSpend"') < sql.index('"LiteLLM_DailyGlobalSpend"') + global_insert = sql.split('INSERT INTO "LiteLLM_DailyGlobalSpend"', 1)[1] + assert global_insert.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") == 1 + assert "api_key" not in global_insert + assert params.count(0.4) == 1 + assert txns == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("entity_type", "entity_field"), + [ + ("team", "team_id"), + ("org", "organization_id"), + ("tag", "tag"), + ("end_user", "end_user_id"), + ("agent", "agent_id"), + ], +) +async def test_other_entity_flushes_leave_the_global_table_alone(entity_type, entity_field): + """Every entity table sees the same request, so writing the rollup from more than one of + them would count each request once per entity type.""" + prisma_client = _RecordingPrisma() + txn = _entity_txn(entity_field, "e-1", "sk-1") + if entity_type == "tag": + txn["request_id"] = "req-1" + + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=MagicMock(), + daily_spend_transactions={"k": txn}, + entity_type=entity_type, + entity_id_field=entity_field, + ) + + (sql, _params) = prisma_client.db.statements[0] + assert "LiteLLM_DailyGlobalSpend" not in sql + + +@pytest.mark.asyncio +async def test_a_failed_chained_user_flush_keeps_every_transaction_for_retry(): + def raise_outage(): + raise ValueError("simulated database outage") + + prisma_client = _RecordingPrisma(execute_raw=raise_outage) + txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(3)} + expected = dict(txns) + mock_proxy_logging = MagicMock() + mock_proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(ValueError, match="simulated database outage"): + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=mock_proxy_logging, + daily_spend_transactions=txns, + entity_type="user", + entity_id_field="user_id", + ) + + assert txns == expected + assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in prisma_client.db.statements[0][0] + + @pytest.mark.asyncio async def test_commit_key_spend_updates_includes_last_active(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 5ff3f89343b..5c74facae6a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -13,7 +13,13 @@ from pytest_postgresql import factories from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR -from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT +import pathlib + +from litellm.constants import ( + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, + PTU_SENTINEL_API_KEY, + USAGE_TOP_API_KEYS_LIMIT, +) from litellm.proxy.management_endpoints.common_daily_activity import ( _adjust_dates_for_timezone, _build_aggregated_sql_query, @@ -23,8 +29,11 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, + key_free_source_table, update_metrics, ) +from litellm.proxy.spend_tracking.daily_global_spend_rollup import RECONCILE_DAY_SQL +from litellm.proxy.utils import evict_config_param from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, SpendMetrics, @@ -1618,6 +1627,149 @@ async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == {"key-1"} +def _prisma_with_marker(marker: str | None) -> MagicMock: + prisma = MagicMock() + prisma.db = MagicMock() + prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + row = None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') + prisma.get_generic_data = AsyncMock(return_value=row) + return prisma + + +def _unfiltered_user_query(**overrides): + return { + "table_name": "litellm_dailyuserspend", + "entity_id_field": "user_id", + "entity_id": None, + "start_date": "2026-06-01", + "end_date": "2026-06-02", + "model": None, + "api_key": None, + "exclude_entity_ids": None, + "timezone_offset_minutes": None, + "include_current_utc_day": False, + **overrides, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("marker", "overrides", "expected"), + [ + ("2026-06-02", {}, "LiteLLM_DailyGlobalSpend"), + ("2026-06-02", {"model": "gpt-5"}, "LiteLLM_DailyGlobalSpend"), + ("2026-06-01", {}, None), + (None, {}, None), + ("2026-06-02", {"api_key": "sk-1"}, None), + ("2026-06-02", {"api_key": []}, None), + ("2026-06-02", {"entity_id": "u-1"}, None), + ("2026-06-02", {"exclude_entity_ids": ["u-1"]}, None), + ("2026-06-02", {"table_name": "litellm_dailyteamspend", "entity_id_field": "team_id"}, None), + ], +) +async def test_key_free_source_table_routes_only_unfiltered_user_reads_within_the_marker(marker, overrides, expected): + """Anything that filters by key or entity has no counterpart in the global table, and a + range the reconcile has not reached must stay on the per-key table.""" + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + prisma = _prisma_with_marker(marker) + + assert await key_free_source_table(prisma, _unfiltered_user_query(**overrides)) == expected + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + +@pytest.mark.asyncio +async def test_key_free_source_table_judges_the_timezone_extended_end_not_the_requested_one(): + """A caller west of UTC asking through their local today gets today's UTC bucket added to + the range; the marker must cover that extended day, not just the requested end.""" + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + today_utc: Final = datetime.now(timezone.utc).date() + yesterday: Final = (today_utc - timedelta(days=1)).isoformat() + query: Final = _unfiltered_user_query( + start_date=yesterday, end_date=yesterday, timezone_offset_minutes=24 * 60, include_current_utc_day=True + ) + + assert await key_free_source_table(_prisma_with_marker(yesterday), query) is None + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + assert await key_free_source_table(_prisma_with_marker(today_utc.isoformat()), query) == "LiteLLM_DailyGlobalSpend" + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + +_GLOBAL_SPEND_MIGRATION: Final = ( + pathlib.Path(__file__).resolve().parents[4] + / "litellm-proxy-extras" + / "litellm_proxy_extras" + / "migrations" + / "20260915000000_add_daily_global_spend" + / "migration.sql" +) + + +@pytest.mark.asyncio +async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_free_arm( + _aggregated_postgresql: psycopg.Connection, +): + """With the range reconciled, the key-free arm reads LiteLLM_DailyGlobalSpend while the + per-key arm stays on the user table, and the response is identical to the all-per-key + read: same totals, same rollups, same top keys.""" + n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 + rows: Final = [ + ( + f"row-{day}-{i:03d}", + f"user-{i % 7}", + day, + f"key-{i:03d}", + "gpt-5" if i % 2 else "claude", + "" if i % 3 else "gpt-5", + "openai" if i % 2 else None, + "/v1/chat/completions" if i % 5 else None, + 10, + float(i + 1), + 1, + 1, + ) + for day in ("2026-06-01", "2026-06-02") + for i in range(n_keys) + ] + _seed_daily_user_spend(_aggregated_postgresql, rows) + with _aggregated_postgresql.cursor() as cur: + cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal + for day in ("2026-06-01", "2026-06-02"): + cur.execute( + re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg + {"p1": day}, + ) + _aggregated_postgresql.commit() + + async def read(marker: str | None, sql_seen: list[str]): + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + prisma = _prisma_with_marker(marker) + run_query = _psycopg_query_raw(_aggregated_postgresql, []) + + async def query_raw(sql: str, *params: str): + sql_seen.append(sql) + return await run_query(sql, *params) + + prisma.db.query_raw = query_raw + return await get_daily_activity_aggregated( + prisma_client=prisma, + entity_metadata_field=None, + **_unfiltered_user_query(), + ) + + per_key_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim + global_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim + from_per_key = await read(None, per_key_sql) + from_global = await read("2026-06-02", global_sql) + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + assert per_key_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 0 + assert global_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 1 + assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 2 + assert from_global.model_dump() == from_per_key.model_dump() + assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT + assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} def _no_spend_record(): """A rollup row for a key with no spend, where SUM() returns NULL (None).""" return SimpleNamespace( diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index deb7289d2d1..ee72e98ffa9 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -1042,6 +1042,54 @@ async def test_spend_report_locks_are_never_released(): proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited() +def _init_daily_global_spend_reconcile_job() -> tuple[MagicMock, MagicMock, MagicMock]: + scheduler = MagicMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.alerting_handler = AsyncMock() + prisma_client = MagicMock() + ProxyStartupEvent._initialize_daily_global_spend_reconcile_job( + scheduler=scheduler, + proxy_logging_obj=proxy_logging_obj, + prisma_client=prisma_client, + ) + return scheduler, proxy_logging_obj, prisma_client + + +def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run(): + """Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a + fresh deploy switches usage reads to the global table without waiting for the nightly + run, and replaces any previous registration of the same job id.""" + from datetime import datetime, timedelta, timezone + + from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID + + scheduler, _, _ = _init_daily_global_spend_reconcile_job() + + (call,) = scheduler.add_job.call_args_list + assert call.kwargs["id"] == DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID + assert call.kwargs["replace_existing"] is True + assert call.args[1:] == ("cron",) + assert (call.kwargs["hour"], call.kwargs["minute"], call.kwargs["timezone"]) == (0, 30, "UTC") + assert timedelta(0) < call.kwargs["next_run_time"] - datetime.now(timezone.utc) <= timedelta(minutes=2) + + +@pytest.mark.asyncio +async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch): + scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job() + run = AsyncMock() + monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run) + + await scheduler.add_job.call_args.args[0]() + + run.assert_awaited_once() + assert run.await_args.args == (prisma_client,) + assert run.await_args.kwargs["pod_lock_manager"] is proxy_logging_obj.db_spend_update_writer.pod_lock_manager + await run.await_args.kwargs["alert"]("day 2026-09-01 failed") + proxy_logging_obj.alerting_handler.assert_awaited_once() + assert proxy_logging_obj.alerting_handler.await_args.kwargs["message"] == "day 2026-09-01 failed" + assert proxy_logging_obj.alerting_handler.await_args.kwargs["level"] == "High" + + @pytest.mark.asyncio async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch): """The boot-time send goes through the same gate, so a losing pod sends nothing at all: diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py new file mode 100644 index 00000000000..13dc757cbbd --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -0,0 +1,382 @@ +"""Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818).""" + +import pathlib +import re +from contextlib import asynccontextmanager +from datetime import date +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import psycopg +import pytest +from psycopg.rows import dict_row +from pytest_postgresql import factories + +from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM +from litellm.proxy.db.daily_spend_bulk_upsert import ( + DAILY_SPEND_TABLES, + build_bulk_upsert_with_global_rollup, + merge_by_conflict_key, +) +from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( + RECONCILE_DAY_SQL, + reconciled_through, + run_daily_global_spend_reconcile, + run_scheduled_daily_global_spend_reconcile, +) +from litellm.proxy.utils import evict_config_param + +USER_TABLE: Final = DAILY_SPEND_TABLES["user"] +TODAY: Final = date(2026, 9, 15) + + +class _FakeConfigRow: + def __init__(self, param_name: str, param_value: object) -> None: + self.param_name = param_name + self.param_value = param_value + + +class _FakeConfigTable: + def __init__(self) -> None: + self.rows: dict[str, object] = {} + + async def upsert(self, *, where: dict[str, str], data: dict[str, dict[str, str]]) -> _FakeConfigRow: + self.rows[where["param_name"]] = data["update"]["param_value"] + return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) + + +class _FakeTransaction: + def __init__(self, prisma: "_FakePrisma") -> None: + self._prisma = prisma + + async def execute_raw(self, sql: str, *params: str) -> int: + if "LOCK TABLE" in sql: + self._prisma.locks_taken += 1 + return 0 + (day,) = params + if day in self._prisma.failing_days: + raise RuntimeError(f"day {day} exploded") + self._prisma.reconciled.append(day) + return 1 + + +class _FakeDb: + def __init__(self, prisma: "_FakePrisma") -> None: + self._prisma = prisma + self.litellm_config = _FakeConfigTable() + + async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: + first, last = params + return [{"date": d} for d in sorted(self._prisma.user_days) if first <= d <= last] + + @asynccontextmanager + async def tx(self, timeout: object): + yield _FakeTransaction(self._prisma) + + +class _FakePrisma: + """Enough of PrismaClient for the reconcile: per-key dates, a config table, and a transaction.""" + + def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: + self.user_days = user_days + self.failing_days = failing_days + self.reconciled: list[str] = [] + self.locks_taken = 0 + self.db = _FakeDb(self) + + async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: + stored = self.db.litellm_config.rows.get(value) + return None if stored is None else _FakeConfigRow(value, stored) + + +@pytest.fixture(autouse=True) +async def _fresh_marker_cache(): + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + yield + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + +@pytest.mark.asyncio +async def test_first_run_rolls_up_every_historical_day_and_today_then_marks_today(): + """Before any marker exists, every day with per-key rows is rolled up, plus today even + with no rows yet, so reads for ranges ending today can switch to the global table.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14")) + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15") + assert result.failed_day is None + assert result.reconciled_through == "2026-09-15" + assert await reconciled_through(prisma) == "2026-09-15" + assert prisma.locks_taken == 4 + + +@pytest.mark.asyncio +async def test_later_run_replays_the_marker_day_and_the_day_before_only(): + """Days older than marker-1 are settled; the marker day and its predecessor are replayed so + rows a pre-writer pod flushed around midnight during a rolling deploy get folded in.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14", "2026-09-15") + assert "2026-09-01" not in prisma.reconciled + assert await reconciled_through(prisma) == "2026-09-15" + + +@pytest.mark.asyncio +async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_good_day(): + """The marker may never claim a day that was not rewritten: reads past it would then trust + a global table missing that day's spend.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01",) + assert result.failed_day == "2026-09-02" + assert result.reconciled_through == "2026-09-01" + assert prisma.reconciled == ["2026-09-01"] + assert await reconciled_through(prisma) == "2026-09-01" + + +@pytest.mark.asyncio +async def test_the_next_run_resumes_from_the_failed_day(): + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) + await run_daily_global_spend_reconcile(prisma, today=TODAY) + prisma.failing_days = frozenset() + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03", "2026-09-15") + assert await reconciled_through(prisma) == "2026-09-15" + + +@pytest.mark.asyncio +async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): + """A pre-writer pod flushing rows for the day before the marker is exactly the replay case; + when that replay fails the marker must stay put and the operator must hear about it.""" + prisma = _FakePrisma(user_days=("2026-09-13",)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + prisma.user_days = ("2026-09-12", "2026-09-13") + prisma.failing_days = frozenset({"2026-09-12"}) + alert = AsyncMock() + + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert, today=TODAY) + + assert result is not None + assert result.days_reconciled == () + assert result.failed_day == "2026-09-12" + assert result.reconciled_through == "2026-09-13" + alert.assert_awaited_once() + assert "2026-09-12" in alert.await_args.args[0] + + +@pytest.mark.asyncio +async def test_a_clean_run_does_not_alert(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + alert = AsyncMock() + + await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert, today=TODAY) + + alert.assert_not_awaited() + + +def _pod_lock(acquired: bool) -> MagicMock: + lock = MagicMock() + lock.redis_cache = MagicMock() + lock.redis_cache.async_get_cache = AsyncMock(return_value="other-pod") + lock.get_redis_lock_key = MagicMock(return_value="lock-key") + lock.acquire_lock = AsyncMock(return_value=acquired) + lock.release_lock = AsyncMock() + return lock + + +@pytest.mark.asyncio +async def test_scheduled_run_skips_when_another_pod_holds_the_lock(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + lock = _pod_lock(acquired=False) + + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + + assert result is None + assert prisma.reconciled == [] + lock.release_lock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + lock = _pod_lock(acquired=True) + + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + + assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + lock.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read(): + """A Redis outage must not stall the backfill: the day rewrite is idempotent, so running + twice is only wasted effort while skipping forever leaves usage on the slow path.""" + prisma = _FakePrisma(user_days=("2026-09-13",)) + lock = _pod_lock(acquired=False) + lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) + + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + + assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + lock.release_lock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_marker_is_read_back_from_the_json_string_the_config_table_stores(): + prisma = _FakePrisma(user_days=()) + prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-10"}' + + assert await reconciled_through(prisma) == "2026-09-10" + + +@pytest.mark.asyncio +async def test_an_unparseable_marker_reads_as_never_reconciled(): + prisma = _FakePrisma(user_days=()) + prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"something_else": 1}' + + assert await reconciled_through(prisma) is None + + +_rollup_postgresql_proc: Final = factories.postgresql_proc() +_rollup_postgresql: Final = factories.postgresql("_rollup_postgresql_proc") + +_MIGRATIONS_DIR: Final = ( + pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" +) +_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" + +_DAILY_USER_SPEND_DDL: Final = """ + CREATE TABLE "LiteLLM_DailyUserSpend" ( + id TEXT PRIMARY KEY, + user_id TEXT, + date TEXT NOT NULL, + api_key TEXT NOT NULL, + model TEXT, + model_group TEXT, + custom_llm_provider TEXT, + mcp_namespaced_tool_name TEXT, + endpoint TEXT, + prompt_tokens BIGINT DEFAULT 0, + completion_tokens BIGINT DEFAULT 0, + cache_read_input_tokens BIGINT DEFAULT 0, + cache_creation_input_tokens BIGINT DEFAULT 0, + compression_saved_tokens BIGINT DEFAULT 0, + compression_savings_spend DOUBLE PRECISION DEFAULT 0, + prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, + spend DOUBLE PRECISION DEFAULT 0, + api_requests BIGINT DEFAULT 0, + successful_requests BIGINT DEFAULT 0, + failed_requests BIGINT DEFAULT 0, + created_at TIMESTAMP DEFAULT now(), + updated_at TIMESTAMP, + UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) + ) +""" + +_PER_KEY_SUMS_SQL: Final = """ + SELECT COALESCE(model, '') AS model, COALESCE(model_group, '') AS model_group, + COALESCE(custom_llm_provider, '') AS custom_llm_provider, + SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests + FROM "LiteLLM_DailyUserSpend" WHERE date = %s + GROUP BY 1, 2, 3 ORDER BY 1, 2, 3 +""" +_GLOBAL_ROWS_SQL: Final = """ + SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests + FROM "LiteLLM_DailyGlobalSpend" WHERE date = %s ORDER BY 1, 2, 3 +""" + + +def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: + converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) + conn.execute( + converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query + {f"p{i}": v for i, v in enumerate(params, start=1)}, + ) + conn.commit() + + +def _user_txn(**overrides): + return { + "user_id": "u-1", + "date": "2026-09-14", + "api_key": "sk-1", + "model": "gpt-5", + "model_group": "gpt-5", + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": "", + "endpoint": "/chat/completions", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 1.0, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + **overrides, + } + + +def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: + return [ + ( + r["model"], + r["model_group"], + r["custom_llm_provider"], + float(r["spend"]), + int(r["prompt_tokens"]), + int(r["api_requests"]), + ) # pyright: ignore[reportArgumentType] # dict_row values are untyped + for r in rows + ] + + +def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_postgresql: psycopg.Connection): + """Against real Postgres and the shipped migration: rows the writer never saw (a + pre-writer pod's flush, NULL and '' dimension spellings) end up folded into the global + day, running the day twice changes nothing, and other days are left alone.""" + conn: Final = _rollup_postgresql + conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal + conn.commit() + + written_batch = merge_by_conflict_key( + USER_TABLE, + (_user_txn(api_key="sk-1", spend=1.0), _user_txn(api_key="sk-2", user_id="u-2", spend=2.0, prompt_tokens=20)), + ) + _execute_dollar_sql(conn, *build_bulk_upsert_with_global_rollup(USER_TABLE, written_batch)) + + conn.execute( + """ + INSERT INTO "LiteLLM_DailyUserSpend" + (id, user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, + endpoint, prompt_tokens, spend, api_requests) + VALUES + ('legacy-1', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', NULL, 'openai', NULL, NULL, 5, 4.0, 1), + ('legacy-2', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', '', 'openai', '', '', 5, 8.0, 1), + ('legacy-3', 'u-9', '2026-09-13', 'sk-9', 'claude', '', 'anthropic', '', '', 7, 16.0, 1) + """ + ) + conn.commit() + + _execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",)) + _execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",)) + + with conn.cursor(row_factory=dict_row) as cur: + global_rows = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-14",)).fetchall() + per_key = cur.execute(_PER_KEY_SUMS_SQL, ("2026-09-14",)).fetchall() + untouched = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-13",)).fetchall() + + assert _normalized(global_rows) == _normalized(per_key) + assert sum(float(r["spend"]) for r in global_rows) == pytest.approx(15.0) # pyright: ignore[reportArgumentType] # dict_row values are untyped + assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")] + assert untouched == [] From ad8de0e1927c18d5d14c92939bbd531c54573874 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:42:19 +0000 Subject: [PATCH 019/102] fix(proxy): roll up only closed days into LiteLLM_DailyGlobalSpend and split the key-free read at the marker The write path no longer dual-writes the global table. The cron rolls up closed UTC days only, so a pod still flushing the current day can never leave the global table short. The key-free arm reads days through the marker from the global table and later days from LiteLLM_DailyUserSpend in one UNION ALL, and the marker comes from the config cache rather than a per-request database lookup. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/daily_spend_bulk_upsert.py | 98 +++--------- litellm/proxy/db/db_spend_update_writer.py | 7 +- .../common_daily_activity.py | 79 ++++++--- .../daily_global_spend_rollup.py | 43 ++--- .../proxy/db/test_daily_spend_bulk_upsert.py | 150 ------------------ .../proxy/db/test_db_spend_update_writer.py | 101 +----------- .../test_common_daily_activity.py | 90 ++++++----- .../test_daily_global_spend_rollup.py | 92 +++++------ 8 files changed, 204 insertions(+), 456 deletions(-) diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index c83043101eb..a143643577e 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -25,41 +25,29 @@ SpendRow = Mapping[str, object] @dataclass(frozen=True, slots=True) class DailySpendTable: - """A daily rollup table and the unique constraint its upserts arbitrate on.""" + """The physical table behind one entity's daily rollup.""" name: str - key_columns: tuple[str, ...] + entity_id_column: str carries_request_id: bool = False +DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( + { + "user": DailySpendTable(name="LiteLLM_DailyUserSpend", entity_id_column="user_id"), + "team": DailySpendTable(name="LiteLLM_DailyTeamSpend", entity_id_column="team_id"), + "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), + "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), + "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), + "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), + } +) + # The unique constraint's columns after the entity id, in constraint order. A NULL can # never match itself in a unique index, so every one of these is normalized to '': the # conflict target has to be NULL-free or the row is re-inserted on every single flush. _KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") - -def _entity_table(name: str, entity_id_column: str, carries_request_id: bool = False) -> DailySpendTable: - return DailySpendTable( - name=name, key_columns=(entity_id_column, *_KEY_COLUMNS), carries_request_id=carries_request_id - ) - - -DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( - { - "user": _entity_table("LiteLLM_DailyUserSpend", "user_id"), - "team": _entity_table("LiteLLM_DailyTeamSpend", "team_id"), - "org": _entity_table("LiteLLM_DailyOrganizationSpend", "organization_id"), - "end_user": _entity_table("LiteLLM_DailyEndUserSpend", "end_user_id"), - "agent": _entity_table("LiteLLM_DailyAgentSpend", "agent_id"), - "tag": _entity_table("LiteLLM_DailyTagSpend", "tag", carries_request_id=True), - } -) - -GLOBAL_SPEND_TABLE: Final = DailySpendTable( - name="LiteLLM_DailyGlobalSpend", - key_columns=("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"), -) - _COUNTER_COLUMNS: Final = ( "prompt_tokens", "completion_tokens", @@ -104,7 +92,7 @@ def _as_float(value: object) -> float: def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" - return tuple(_as_text(transaction.get(column)) for column in table.key_columns) + return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -142,11 +130,7 @@ def _row_params( return ( str(uuid.uuid4()), *key, - *( - () - if "model_group" in table.key_columns - else (None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")),) - ), + None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")), *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), @@ -156,25 +140,26 @@ def _row_params( def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - *table.key_columns, - *(() if "model_group" in table.key_columns else ("model_group",)), + table.entity_id_column, + *_KEY_COLUMNS, + "model_group", *_COUNTER_COLUMNS, *_SPEND_COLUMNS, *(("request_id",) if table.carries_request_id else ()), ) -def _upsert_statement( +def build_bulk_upsert( table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]], - first_param: int, -) -> str: +) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one merged batch, plus its positional arguments.""" columns: Final = _insert_columns(table) quoted_table: Final = f'"{table.name}"' rows: Final = ", ".join( "(" + ", ".join( - f"${first_param + row_index * len(columns) + offset}::{_CASTS.get(column, 'text')}" + f"${row_index * len(columns) + offset + 1}::{_CASTS.get(column, 'text')}" for offset, column in enumerate(columns) ) + ", (NOW() AT TIME ZONE 'UTC'))" @@ -191,44 +176,11 @@ def _upsert_statement( if table.carries_request_id else "" ) - return ( + sql: Final = ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted(table.key_columns)}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) - - -def _params(table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]]) -> tuple[SqlValue, ...]: - return tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) - - -def build_bulk_upsert( - table: DailySpendTable, - batch: Sequence[tuple[tuple[str, ...], SpendRow]], -) -> tuple[str, tuple[SqlValue, ...]]: - """The single statement writing one merged batch, plus its positional arguments.""" - return _upsert_statement(table, batch, first_param=1), _params(table, batch) - - -def build_bulk_upsert_with_global_rollup( - table: DailySpendTable, - batch: Sequence[tuple[tuple[str, ...], SpendRow]], -) -> tuple[str, tuple[SqlValue, ...]]: - """One statement writing a batch to its table and, atomically, its key-free rollup - to ``LiteLLM_DailyGlobalSpend``. - - A data-modifying CTE runs both inserts in the same snapshot and transaction, so a - batch that lands in one table lands in both and a retried deadlock replays both. - Postgres does not order the CTE against the main statement, so two writers can still - deadlock across the tables; the caller's deadlock retry covers that, and each insert - takes its own rows in key order so same-table lock order stays deterministic. - """ - global_batch: Final = merge_by_conflict_key(GLOBAL_SPEND_TABLE, tuple(row for _, row in batch)) - entity_params: Final = _params(table, batch) - sql: Final = ( - f"WITH entity_rows AS (\n{_upsert_statement(table, batch, first_param=1)}\nRETURNING 1)\n" - f"{_upsert_statement(GLOBAL_SPEND_TABLE, global_batch, first_param=len(entity_params) + 1)}" - ) - return sql, (*entity_params, *_params(GLOBAL_SPEND_TABLE, global_batch)) + return sql, tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d5c839a9be8..eaa03c5d7f7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -46,7 +46,6 @@ from litellm.proxy._types import ( from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, - build_bulk_upsert_with_global_rollup, merge_by_conflict_key, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -1940,11 +1939,7 @@ class DBSpendUpdateWriter: merged_batch = merge_by_conflict_key( table=table, transactions=tuple(transactions_to_process.values()) ) - sql, params = ( - build_bulk_upsert_with_global_rollup(table=table, batch=merged_batch) - if entity_type == "user" - else build_bulk_upsert(table=table, batch=merged_batch) - ) + sql, params = build_bulk_upsert(table=table, batch=merged_batch) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: # Log detailed error information for debugging batch upsert failures diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index f1d78dca201..b90f874c04c 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -11,8 +11,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT from litellm.proxy._types import CommonProxyErrors -from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE -from litellm.proxy.spend_tracking.daily_global_spend_rollup import reconciled_through +from litellm.proxy.spend_tracking.daily_global_spend_rollup import GLOBAL_SPEND_TABLE_NAME, reconciled_through from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, recover_double_hashed_key_metadata, @@ -736,28 +735,62 @@ def _rollup_metric_select(table_name: str) -> str: _MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" -async def key_free_source_table(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: - """The table the key-free arm reads from, when the global rollup can answer instead of the per-key table. +_KEY_FREE_SOURCE_COLUMNS: Final = ( + "date", + "model", + "model_group", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + "spend", + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "compression_saved_tokens", + "compression_savings_spend", + "prompt_caching_savings_spend", + "gateway_injected_caching_savings_spend", + "autorouter_savings_spend", + "api_requests", + "successful_requests", + "failed_requests", +) - Only an unfiltered read of the user table has the same rows as ``LiteLLM_DailyGlobalSpend``, - and only through the day the reconcile marker has reached: the writer keeps that day - current, later days are covered once the next run advances the marker. + +async def global_rollup_reconciled_through(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: + """The last day ``LiteLLM_DailyGlobalSpend`` can answer the key-free arm for, or None to + read it all from the per-key table. + + Only an unfiltered read of the user table sums to the same rows as the global table. The + marker read is served from the config cache, so this is not a database round trip per request. """ if query["table_name"] != "litellm_dailyuserspend": return None if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]: return None - _, adjusted_end = _adjust_dates_for_timezone( - query["start_date"], query["end_date"], query["timezone_offset_minutes"], query["include_current_utc_day"] - ) try: - marker: Final = await reconciled_through(prisma_client) + return await reconciled_through(prisma_client) except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc) return None - if marker is None or adjusted_end > marker: - return None - return GLOBAL_SPEND_TABLE.name + + +def _key_free_source(pg_table: str, where_clause: str, marker_param: str | None) -> str: + """The relation the key-free arm aggregates: the per-key table alone, or the global rollup + for days through the marker plus the per-key table for the days still open after it.""" + if marker_param is None: + return f'"{pg_table}"\n WHERE {where_clause}' + columns: Final = ", ".join(_KEY_FREE_SOURCE_COLUMNS) + return f"""( + SELECT {columns} + FROM "{GLOBAL_SPEND_TABLE_NAME}" + WHERE {where_clause} AND date <= {marker_param} + UNION ALL + SELECT {columns} + FROM "{pg_table}" + WHERE {where_clause} AND date > {marker_param} + ) AS key_free_source""" def _build_aggregated_sql_query( @@ -772,15 +805,16 @@ def _build_aggregated_sql_query( exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path timezone_offset_minutes: int | None = None, include_current_utc_day: bool = False, - key_free_table: str | None = None, + global_rollup_through: str | None = None, ) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params """Build the GROUPING SETS query for aggregated daily activity. One statement, two UNION ALL arms over the same WHERE clause. The first arm is key-free: grand total, per-date totals and the (date, model / model_group / provider / mcp / endpoint) rollups, so its row count never grows with the number - of keys; it reads ``key_free_table`` when given (the global rollup, whose row count - never grew with the number of keys to begin with) and the entity table otherwise. + of keys. With ``global_rollup_through`` it reads days through that marker from + ``LiteLLM_DailyGlobalSpend`` (whose row count never grew with the number of keys to + begin with) and only the days after it from the per-key table. The second arm emits the (date, , api_key) rollups for the USAGE_TOP_API_KEYS_LIMIT highest-spend keys only. Both arms share the 7-bit group_level bitmask (date, api_key, model, model_group, provider, mcp, endpoint). @@ -806,8 +840,8 @@ def _build_aggregated_sql_query( exclude_entity_ids=exclude_entity_ids, ) sentinel_param: Final = f"${len(where_params) + 1}" + marker_param: Final = None if global_rollup_through is None else f"${len(where_params) + 2}" metric_select: Final = _rollup_metric_select(table_name) - key_free_source: Final = key_free_table or pg_table # TODO: drop the successful_requests/failed_requests aggregates (and the # total_successful_requests metadata they feed) once the admin UI reads SGR @@ -826,8 +860,7 @@ def _build_aggregated_sql_query( | GROUPING(model, {_MODEL_GROUP_EXPR}, custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level,{metric_select} - FROM "{key_free_source}" - WHERE {where_clause} + FROM {_key_free_source(pg_table, where_clause, marker_param)} GROUP BY GROUPING SETS ( (date), (date, model), @@ -869,7 +902,8 @@ def _build_aggregated_sql_query( )) """ - return sql_query, [*where_params, PTU_SENTINEL_API_KEY] + marker_params: Final = () if global_rollup_through is None else (global_rollup_through,) + return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params] def _build_entity_rollup_sql_query( @@ -1418,7 +1452,8 @@ async def get_daily_activity_aggregated( include_current_utc_day=include_current_utc_day, ) sql_query, sql_params = _build_aggregated_sql_query( - **query_kwargs, key_free_table=await key_free_source_table(prisma_client, query_kwargs) + **query_kwargs, + global_rollup_through=await global_rollup_reconciled_through(prisma_client, query_kwargs), ) entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index 9d344421332..a9fb7669785 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -1,9 +1,10 @@ -"""Reconcile ``LiteLLM_DailyGlobalSpend`` from ``LiteLLM_DailyUserSpend``, one day per transaction. +"""Roll closed UTC days of ``LiteLLM_DailyUserSpend`` up into ``LiteLLM_DailyGlobalSpend``. -The spend writer keeps both tables in step from the moment it is deployed; this job rolls up -the days before that and records how far it has reached in ``LiteLLM_Config`` so usage reads -know when the global table can answer for a date range. It runs as a background cron, never -in a Prisma migration, since on a large deployment the aggregate is minutes of work. +Only days that are over get rolled up, so a pod still flushing per-key spend for the current +day can never leave the global table short; usage reads serve days through the recorded +marker from the global table and later days live from the per-key table. The marker lives in +``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on a +large deployment the first backfill is minutes of work. """ from collections.abc import Awaitable, Callable @@ -19,7 +20,6 @@ from litellm.constants import ( DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ) -from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE from litellm.repositories.config_repository import ConfigRepository if TYPE_CHECKING: @@ -27,8 +27,11 @@ if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient -_DAY_TRANSACTION_TIMEOUT: Final = timedelta(minutes=10) _REPLAY_DAYS: Final = 1 +GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend" +# The unique constraint, in constraint order. NULL never matches itself in a unique index, so +# every column is normalized to '' or the same group would be inserted again on every run. +_KEY_COLUMNS: Final = ("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") _METRIC_COLUMNS: Final = ( "prompt_tokens", "completion_tokens", @@ -51,23 +54,21 @@ def _quoted(columns: tuple[str, ...]) -> str: def _reconcile_day_sql() -> str: - key_columns: Final = GLOBAL_SPEND_TABLE.key_columns - normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in key_columns) + normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in _KEY_COLUMNS) sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS) overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS) return ( - f'INSERT INTO "{GLOBAL_SPEND_TABLE.name}" ("id", {_quoted(key_columns)}, {_quoted(_METRIC_COLUMNS)}, ' + f'INSERT INTO "{GLOBAL_SPEND_TABLE_NAME}" ("id", {_quoted(_KEY_COLUMNS)}, {_quoted(_METRIC_COLUMNS)}, ' '"updated_at")\n' f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n" 'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n' f"GROUP BY {normalized_keys}\n" - f"ON CONFLICT ({_quoted(key_columns)}) DO UPDATE SET {overwrite}, " + f"ON CONFLICT ({_quoted(_KEY_COLUMNS)}) DO UPDATE SET {overwrite}, " "\"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) RECONCILE_DAY_SQL: Final = _reconcile_day_sql() -_LOCK_GLOBAL_TABLE_SQL: Final = f'LOCK TABLE "{GLOBAL_SPEND_TABLE.name}" IN EXCLUSIVE MODE' _PENDING_DAYS_SQL: Final = ( 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" >= $1 AND "date" <= $2 ORDER BY "date"' ) @@ -134,19 +135,19 @@ def _first_pending_day(marker: str | None) -> str: async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: - """Every UTC day through today still to roll up, oldest first; the marker day and the one - before it are replayed so rows flushed by a pre-writer pod during a rolling deploy are folded in.""" + """Every closed UTC day (strictly before today) still to roll up, oldest first. The marker + day and the one before it are replayed so per-key rows that landed after their day was + rolled up (a flush straddling midnight, a late retry) are folded in.""" marker: Final = await reconciled_through(prisma_client) - rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), today.isoformat()) - return tuple(sorted({*(_DateRow.model_validate(row).date for row in rows), today.isoformat()})) + last_closed_day: Final = (today - timedelta(days=1)).isoformat() + rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), last_closed_day) + return tuple(_DateRow.model_validate(row).date for row in rows) async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: - """Rewrite one day of the global table from the per-key sums; the table lock keeps the - writer's increments out between the aggregate and the overwrite so none are lost.""" - async with prisma_client.db.tx(timeout=_DAY_TRANSACTION_TIMEOUT) as transaction: - await transaction.execute_raw(_LOCK_GLOBAL_TABLE_SQL) - await transaction.execute_raw(RECONCILE_DAY_SQL, day) + """Rewrite one day of the global table from the per-key sums. Idempotent: a rerun + overwrites every group with the same totals.""" + await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day) async def run_daily_global_spend_reconcile( diff --git a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py index cc443a2cfe5..c1efb3e7220 100644 --- a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py @@ -1,19 +1,12 @@ """Tests for the single-statement daily spend upsert (LIT-5291).""" -import pathlib import re -from typing import Final -import psycopg import pytest -from psycopg.rows import dict_row -from pytest_postgresql import factories from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, - GLOBAL_SPEND_TABLE, build_bulk_upsert, - build_bulk_upsert_with_global_rollup, conflict_key, merge_by_conflict_key, ) @@ -192,146 +185,3 @@ async def test_writer_survives_a_transaction_whose_key_columns_are_null(): _, params = prisma_client.db.statements[0] assert None not in params[:9] assert transactions == {} - - -def user_txn(**overrides): - txn = {**tag_txn(), "user_id": "u-1", **overrides} - del txn["tag"] - del txn["request_id"] - return txn - - -def _bound_rows(insert_sql: str, params: tuple[object, ...]) -> list[dict[str, object]]: - """Each VALUES row of one INSERT as a column -> bound value mapping, consuming params in order.""" - header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", insert_sql) - assert header is not None, insert_sql - columns = [c.strip('"') for c in header.group(1).split(", ") if c != '"updated_at"'] - row_count = insert_sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") - return [dict(zip(columns, params[i * len(columns) : (i + 1) * len(columns)])) for i in range(row_count)] - - -def test_global_rollup_folds_every_key_and_user_into_one_row_per_dimension_tuple(): - """The global table has no api_key or user_id, so a batch spread over many keys and - users must collapse to one row per (date, model, group, provider, mcp, endpoint).""" - batch = merge_by_conflict_key( - USER_TABLE, - tuple(user_txn(user_id=f"u-{i}", api_key=f"sk-{i}", spend=1.0, api_requests=1) for i in range(5)) - + (user_txn(user_id="u-0", api_key="sk-0", model="claude", spend=10.0, api_requests=3),), - ) - - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - - entity_insert, global_insert = sql.split("RETURNING 1)") - entity_rows = _bound_rows(entity_insert, params) - global_rows = _bound_rows(global_insert, params[len(entity_rows) * len(entity_rows[0]) :]) - assert len(entity_rows) == 6 - assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in global_insert - assert [(r["model"], r["spend"], r["api_requests"]) for r in global_rows] == [ - ("claude", 10.0, 3), - ("gpt-4o-mini", 5.0, 5), - ] - assert all("api_key" not in r and "user_id" not in r for r in global_rows) - conflict = re.search(r"ON CONFLICT \(([^)]*)\)", global_insert) - assert conflict is not None - assert conflict.group(1) == ", ".join(f'"{c}"' for c in GLOBAL_SPEND_TABLE.key_columns) - - -def test_global_rollup_params_follow_the_entity_params_in_one_placeholder_sequence(): - """Both inserts bind from one flat tuple, so the global arm's placeholders must start - exactly where the entity arm's stop or every value lands one column off.""" - batch = merge_by_conflict_key(USER_TABLE, (user_txn(),)) - - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - - placeholders = [int(n) for n in re.findall(r"\$(\d+)::", sql)] - assert placeholders == list(range(1, len(params) + 1)) - - -_bulk_upsert_postgresql_proc: Final = factories.postgresql_proc() -_bulk_upsert_postgresql: Final = factories.postgresql("_bulk_upsert_postgresql_proc") - -_MIGRATIONS_DIR: Final = ( - pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" -) -_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" - -_DAILY_USER_SPEND_DDL: Final = """ - CREATE TABLE "LiteLLM_DailyUserSpend" ( - id TEXT PRIMARY KEY, - user_id TEXT, - date TEXT NOT NULL, - api_key TEXT NOT NULL, - model TEXT, - model_group TEXT, - custom_llm_provider TEXT, - mcp_namespaced_tool_name TEXT, - endpoint TEXT, - prompt_tokens BIGINT DEFAULT 0, - completion_tokens BIGINT DEFAULT 0, - cache_read_input_tokens BIGINT DEFAULT 0, - cache_creation_input_tokens BIGINT DEFAULT 0, - compression_saved_tokens BIGINT DEFAULT 0, - compression_savings_spend DOUBLE PRECISION DEFAULT 0, - prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, - spend DOUBLE PRECISION DEFAULT 0, - api_requests BIGINT DEFAULT 0, - successful_requests BIGINT DEFAULT 0, - failed_requests BIGINT DEFAULT 0, - created_at TIMESTAMP DEFAULT now(), - updated_at TIMESTAMP, - UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) - ) -""" - - -def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: - converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) - conn.execute( - converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query - {f"p{i}": v for i, v in enumerate(params, start=1)}, - ) - conn.commit() - - -def test_global_rollup_equals_the_per_key_sums_after_repeated_flushes(_bulk_upsert_postgresql: psycopg.Connection): - """Against real Postgres and the shipped migration: two flushes of a mixed batch leave - the global table exactly equal to the per-key table summed over user and key, with the - NULL and '' spellings of a dimension folded into one row.""" - conn: Final = _bulk_upsert_postgresql - conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal - conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - conn.commit() - - batch = merge_by_conflict_key( - USER_TABLE, - ( - user_txn(user_id="u-1", api_key="sk-1", spend=1.0, prompt_tokens=10), - user_txn(user_id="u-2", api_key="sk-2", spend=2.0, prompt_tokens=20), - user_txn(user_id="u-1", api_key="sk-3", model=None, custom_llm_provider=None, spend=4.0), - user_txn(user_id="u-3", api_key="sk-4", model="", custom_llm_provider="", spend=8.0), - ), - ) - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - _execute_dollar_sql(conn, sql, params) - _execute_dollar_sql(conn, sql, params) - - with conn.cursor(row_factory=dict_row) as cur: - global_rows = cur.execute( - 'SELECT model, spend, prompt_tokens, api_requests FROM "LiteLLM_DailyGlobalSpend" ORDER BY model' - ).fetchall() - per_key = cur.execute( - """ - SELECT COALESCE(model, '') AS model, SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, - SUM(api_requests) AS api_requests - FROM "LiteLLM_DailyUserSpend" GROUP BY COALESCE(model, '') ORDER BY 1 - """ - ).fetchall() - - assert [row["model"] for row in global_rows] == ["", "gpt-4o-mini"] - assert [(r["model"], r["spend"], int(r["prompt_tokens"]), int(r["api_requests"])) for r in global_rows] == [ - (r["model"], float(r["spend"]), int(r["prompt_tokens"]), int(r["api_requests"])) for r in per_key - ] - assert global_rows[0]["spend"] == pytest.approx(24.0) - assert global_rows[1]["spend"] == pytest.approx(6.0) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index d8a9013398e..5e977712a1e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -254,19 +254,14 @@ class _RecordingPrisma: def _row_values(statement: Statement, column: str) -> list[object]: - """Every row's value for one column of the first INSERT, read out of the flat parameter tuple. - - The user-table statement chains a global rollup INSERT after its own, so the row count - comes from the first INSERT's VALUES rather than from the parameter count. - """ + """Every row's value for one column, read out of the flat parameter tuple.""" sql, params = statement header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", sql) assert header is not None, sql columns = header.group(1).split(", ") stride = len(columns) - 1 # updated_at is inlined, not bound offset = columns.index(f'"{column}"') - rows = sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") - return [params[row * stride + offset] for row in range(rows)] + return [params[row * stride + offset] for row in range(len(params) // stride)] @pytest.mark.asyncio @@ -1468,98 +1463,6 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry(): assert daily_spend_transactions == expected -def _entity_txn(entity_field: str, entity_id: str, api_key: str) -> dict[str, object]: - txn = _daily_txn() - del txn["user_id"] - return {**txn, entity_field: entity_id, "api_key": api_key} - - -@pytest.mark.asyncio -async def test_user_flush_writes_the_global_rollup_in_the_same_statement(): - """The user flush is the one place per-key spend becomes key-free spend, so a batch spread - over many keys must land in LiteLLM_DailyGlobalSpend as one row in the same statement. - A separate statement would let a crash between the two leave the tables out of sync.""" - prisma_client = _RecordingPrisma() - txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(4)} - - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=MagicMock(), - daily_spend_transactions=txns, - entity_type="user", - entity_id_field="user_id", - ) - - assert len(prisma_client.db.statements) == 1 - sql, params = prisma_client.db.statements[0] - assert sql.count('INSERT INTO "LiteLLM_DailyUserSpend"') == 1 - assert sql.count('INSERT INTO "LiteLLM_DailyGlobalSpend"') == 1 - assert sql.index('"LiteLLM_DailyUserSpend"') < sql.index('"LiteLLM_DailyGlobalSpend"') - global_insert = sql.split('INSERT INTO "LiteLLM_DailyGlobalSpend"', 1)[1] - assert global_insert.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") == 1 - assert "api_key" not in global_insert - assert params.count(0.4) == 1 - assert txns == {} - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("entity_type", "entity_field"), - [ - ("team", "team_id"), - ("org", "organization_id"), - ("tag", "tag"), - ("end_user", "end_user_id"), - ("agent", "agent_id"), - ], -) -async def test_other_entity_flushes_leave_the_global_table_alone(entity_type, entity_field): - """Every entity table sees the same request, so writing the rollup from more than one of - them would count each request once per entity type.""" - prisma_client = _RecordingPrisma() - txn = _entity_txn(entity_field, "e-1", "sk-1") - if entity_type == "tag": - txn["request_id"] = "req-1" - - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=MagicMock(), - daily_spend_transactions={"k": txn}, - entity_type=entity_type, - entity_id_field=entity_field, - ) - - (sql, _params) = prisma_client.db.statements[0] - assert "LiteLLM_DailyGlobalSpend" not in sql - - -@pytest.mark.asyncio -async def test_a_failed_chained_user_flush_keeps_every_transaction_for_retry(): - def raise_outage(): - raise ValueError("simulated database outage") - - prisma_client = _RecordingPrisma(execute_raw=raise_outage) - txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(3)} - expected = dict(txns) - mock_proxy_logging = MagicMock() - mock_proxy_logging.failure_handler = AsyncMock() - - with pytest.raises(ValueError, match="simulated database outage"): - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=mock_proxy_logging, - daily_spend_transactions=txns, - entity_type="user", - entity_id_field="user_id", - ) - - assert txns == expected - assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in prisma_client.db.statements[0][0] - - @pytest.mark.asyncio async def test_commit_key_spend_updates_includes_last_active(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 5c74facae6a..12e5fe6af4d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1,3 +1,4 @@ +import pathlib import re from collections.abc import Sequence from datetime import datetime, timedelta, timezone @@ -10,11 +11,6 @@ import pytest from psycopg.rows import dict_row from pytest_postgresql import factories -from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR - - -import pathlib - from litellm.constants import ( DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, PTU_SENTINEL_API_KEY, @@ -29,10 +25,11 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, - key_free_source_table, + global_rollup_reconciled_through, update_metrics, ) from litellm.proxy.spend_tracking.daily_global_spend_rollup import RECONCILE_DAY_SQL +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.proxy.utils import evict_config_param from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, @@ -1632,7 +1629,9 @@ def _prisma_with_marker(marker: str | None) -> MagicMock: prisma.db = MagicMock() prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - row = None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') + row = ( + None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') + ) prisma.get_generic_data = AsyncMock(return_value=row) return prisma @@ -1657,9 +1656,9 @@ def _unfiltered_user_query(**overrides): @pytest.mark.parametrize( ("marker", "overrides", "expected"), [ - ("2026-06-02", {}, "LiteLLM_DailyGlobalSpend"), - ("2026-06-02", {"model": "gpt-5"}, "LiteLLM_DailyGlobalSpend"), - ("2026-06-01", {}, None), + ("2026-06-02", {}, "2026-06-02"), + ("2026-06-02", {"model": "gpt-5"}, "2026-06-02"), + ("2026-05-01", {}, "2026-05-01"), (None, {}, None), ("2026-06-02", {"api_key": "sk-1"}, None), ("2026-06-02", {"api_key": []}, None), @@ -1668,33 +1667,50 @@ def _unfiltered_user_query(**overrides): ("2026-06-02", {"table_name": "litellm_dailyteamspend", "entity_id_field": "team_id"}, None), ], ) -async def test_key_free_source_table_routes_only_unfiltered_user_reads_within_the_marker(marker, overrides, expected): - """Anything that filters by key or entity has no counterpart in the global table, and a - range the reconcile has not reached must stay on the per-key table.""" +async def test_global_rollup_marker_is_used_only_for_unfiltered_user_reads(marker, overrides, expected): + """Anything that filters by key or entity has no counterpart in the global table; the + SQL splits the range at the marker itself, so the marker passes through unchanged.""" await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) prisma = _prisma_with_marker(marker) - assert await key_free_source_table(prisma, _unfiltered_user_query(**overrides)) == expected + assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query(**overrides)) == expected await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) @pytest.mark.asyncio -async def test_key_free_source_table_judges_the_timezone_extended_end_not_the_requested_one(): - """A caller west of UTC asking through their local today gets today's UTC bucket added to - the range; the marker must cover that extended day, not just the requested end.""" +async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table(): await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - today_utc: Final = datetime.now(timezone.utc).date() - yesterday: Final = (today_utc - timedelta(days=1)).isoformat() - query: Final = _unfiltered_user_query( - start_date=yesterday, end_date=yesterday, timezone_offset_minutes=24 * 60, include_current_utc_day=True - ) + prisma = _prisma_with_marker(None) + prisma.get_generic_data = AsyncMock(side_effect=RuntimeError("db down")) - assert await key_free_source_table(_prisma_with_marker(yesterday), query) is None - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - assert await key_free_source_table(_prisma_with_marker(today_utc.isoformat()), query) == "LiteLLM_DailyGlobalSpend" + assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query()) is None await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) +def test_aggregated_sql_splits_the_key_free_arm_at_the_marker_and_keeps_the_key_arm_per_key(): + sql, params = _build_aggregated_sql_query(**_unfiltered_user_query(), global_rollup_through="2026-06-01") + marker_param: Final = f"${len(params)}" + + assert params[-1] == "2026-06-01" + assert ( + f'FROM "LiteLLM_DailyGlobalSpend"\n WHERE date >= $1 AND date <= $2 AND date <= {marker_param}' + in sql + ) + assert ( + f'FROM "LiteLLM_DailyUserSpend"\n WHERE date >= $1 AND date <= $2 AND date > {marker_param}' in sql + ) + key_arm: Final = sql.split("UNION ALL\n (WITH top_api_keys")[1] + assert "LiteLLM_DailyGlobalSpend" not in key_arm + assert marker_param not in key_arm + + +def test_aggregated_sql_without_a_marker_reads_the_per_key_table_only(): + sql, params = _build_aggregated_sql_query(**_unfiltered_user_query()) + + assert "LiteLLM_DailyGlobalSpend" not in sql + assert params[-1] == PTU_SENTINEL_API_KEY + + _GLOBAL_SPEND_MIGRATION: Final = ( pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" @@ -1706,12 +1722,12 @@ _GLOBAL_SPEND_MIGRATION: Final = ( @pytest.mark.asyncio -async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_free_arm( +async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_table_and_open_days_live( _aggregated_postgresql: psycopg.Connection, ): - """With the range reconciled, the key-free arm reads LiteLLM_DailyGlobalSpend while the - per-key arm stays on the user table, and the response is identical to the all-per-key - read: same totals, same rollups, same top keys.""" + """Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must + give the same response as reading everything per-key: day 1 from the global table, day 2 + live, one grand total across both. The per-key arm stays on the user table throughout.""" n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 rows: Final = [ ( @@ -1734,11 +1750,10 @@ async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_ _seed_daily_user_spend(_aggregated_postgresql, rows) with _aggregated_postgresql.cursor() as cur: cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - for day in ("2026-06-01", "2026-06-02"): - cur.execute( - re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg - {"p1": day}, - ) + cur.execute( + re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg + {"p1": "2026-06-01"}, + ) _aggregated_postgresql.commit() async def read(marker: str | None, sql_seen: list[str]): @@ -1760,16 +1775,19 @@ async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_ per_key_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim global_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim from_per_key = await read(None, per_key_sql) - from_global = await read("2026-06-02", global_sql) + from_global = await read("2026-06-01", global_sql) await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) assert per_key_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 0 assert global_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 1 - assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 2 + assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 3 assert from_global.model_dump() == from_per_key.model_dump() assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} + + def _no_spend_record(): """A rollup row for a key with no spend, where SUM() returns NULL (None).""" return SimpleNamespace( diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 13dc757cbbd..11ca72e7b3d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -2,7 +2,6 @@ import pathlib import re -from contextlib import asynccontextmanager from datetime import date from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -13,11 +12,7 @@ from psycopg.rows import dict_row from pytest_postgresql import factories from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM -from litellm.proxy.db.daily_spend_bulk_upsert import ( - DAILY_SPEND_TABLES, - build_bulk_upsert_with_global_rollup, - merge_by_conflict_key, -) +from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( RECONCILE_DAY_SQL, reconciled_through, @@ -45,21 +40,6 @@ class _FakeConfigTable: return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) -class _FakeTransaction: - def __init__(self, prisma: "_FakePrisma") -> None: - self._prisma = prisma - - async def execute_raw(self, sql: str, *params: str) -> int: - if "LOCK TABLE" in sql: - self._prisma.locks_taken += 1 - return 0 - (day,) = params - if day in self._prisma.failing_days: - raise RuntimeError(f"day {day} exploded") - self._prisma.reconciled.append(day) - return 1 - - class _FakeDb: def __init__(self, prisma: "_FakePrisma") -> None: self._prisma = prisma @@ -69,19 +49,21 @@ class _FakeDb: first, last = params return [{"date": d} for d in sorted(self._prisma.user_days) if first <= d <= last] - @asynccontextmanager - async def tx(self, timeout: object): - yield _FakeTransaction(self._prisma) + async def execute_raw(self, sql: str, *params: str) -> int: + (day,) = params + if day in self._prisma.failing_days: + raise RuntimeError(f"day {day} exploded") + self._prisma.reconciled.append(day) + return 1 class _FakePrisma: - """Enough of PrismaClient for the reconcile: per-key dates, a config table, and a transaction.""" + """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: self.user_days = user_days self.failing_days = failing_days self.reconciled: list[str] = [] - self.locks_taken = 0 self.db = _FakeDb(self) async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: @@ -97,33 +79,45 @@ async def _fresh_marker_cache(): @pytest.mark.asyncio -async def test_first_run_rolls_up_every_historical_day_and_today_then_marks_today(): - """Before any marker exists, every day with per-key rows is rolled up, plus today even - with no rows yet, so reads for ranges ending today can switch to the global table.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14")) +async def test_first_run_rolls_up_every_closed_day_and_never_today(): + """Before any marker exists every closed day with per-key rows is rolled up. Today is left + out: pods are still flushing it, so it is served live from the per-key table until it closes.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15")) result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15") + assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14") assert result.failed_day is None - assert result.reconciled_through == "2026-09-15" - assert await reconciled_through(prisma) == "2026-09-15" - assert prisma.locks_taken == 4 + assert result.reconciled_through == "2026-09-14" + assert await reconciled_through(prisma) == "2026-09-14" + assert "2026-09-15" not in prisma.reconciled @pytest.mark.asyncio async def test_later_run_replays_the_marker_day_and_the_day_before_only(): """Days older than marker-1 are settled; the marker day and its predecessor are replayed so - rows a pre-writer pod flushed around midnight during a rolling deploy get folded in.""" + per-key rows that landed after their day was rolled up get folded in.""" prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.reconciled.clear() result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14", "2026-09-15") + assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14") assert "2026-09-01" not in prisma.reconciled - assert await reconciled_through(prisma) == "2026-09-15" + assert await reconciled_through(prisma) == "2026-09-14" + + +@pytest.mark.asyncio +async def test_a_run_with_no_new_closed_days_keeps_the_marker(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + + assert result.days_reconciled == ("2026-09-13",) + assert result.reconciled_through == "2026-09-13" @pytest.mark.asyncio @@ -149,16 +143,16 @@ async def test_the_next_run_resumes_from_the_failed_day(): result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03", "2026-09-15") - assert await reconciled_through(prisma) == "2026-09-15" + assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03") + assert await reconciled_through(prisma) == "2026-09-03" @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): - """A pre-writer pod flushing rows for the day before the marker is exactly the replay case; - when that replay fails the marker must stay put and the operator must hear about it.""" + """A late flush for the day before the marker is exactly the replay case; when that replay + fails the marker must stay put and the operator must hear about it.""" prisma = _FakePrisma(user_days=("2026-09-13",)) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.user_days = ("2026-09-12", "2026-09-13") prisma.failing_days = frozenset({"2026-09-12"}) alert = AsyncMock() @@ -212,7 +206,7 @@ async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) - assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_awaited_once() @@ -226,7 +220,7 @@ async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read() result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) - assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_not_awaited() @@ -341,9 +335,9 @@ def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_postgresql: psycopg.Connection): - """Against real Postgres and the shipped migration: rows the writer never saw (a - pre-writer pod's flush, NULL and '' dimension spellings) end up folded into the global - day, running the day twice changes nothing, and other days are left alone.""" + """Against real Postgres and the shipped migration: writer-shaped rows and legacy rows + (NULL and '' dimension spellings) fold into one global day, running the day twice changes + nothing, and other days are left alone.""" conn: Final = _rollup_postgresql conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal @@ -353,7 +347,7 @@ def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_p USER_TABLE, (_user_txn(api_key="sk-1", spend=1.0), _user_txn(api_key="sk-2", user_id="u-2", spend=2.0, prompt_tokens=20)), ) - _execute_dollar_sql(conn, *build_bulk_upsert_with_global_rollup(USER_TABLE, written_batch)) + _execute_dollar_sql(conn, *build_bulk_upsert(USER_TABLE, written_batch)) conn.execute( """ From 84c098df92f8d89ed5d083466ec62347110062e0 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 00:42:03 +0000 Subject: [PATCH 020/102] fix(proxy): fold late-arriving per-key spend into already rolled-up global days The reconcile now records the database clock of the scan behind the last complete run and, on the next run, rewrites every closed day with per-key rows updated since then, however old the day is. Replaying only the marker day and the one before it missed a delayed flush or retry that landed on an older date, and reads through the marker come from the global table alone, so that spend was never counted. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../daily_global_spend_rollup.py | 128 +++++++++++++----- .../test_daily_global_spend_rollup.py | 88 ++++++++++-- 2 files changed, 168 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index a9fb7669785..73068381dab 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -2,9 +2,12 @@ Only days that are over get rolled up, so a pod still flushing per-key spend for the current day can never leave the global table short; usage reads serve days through the recorded -marker from the global table and later days live from the per-key table. The marker lives in -``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on a -large deployment the first backfill is minutes of work. +marker from the global table and later days live from the per-key table. Per-key rows are +dated by request start, so spend can land on a day that was already rolled up (a flush +straddling midnight, a retry after an outage). Each run therefore also rewrites every closed +day that has rows touched since the previous run's scan, whatever the date. The marker lives +in ``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on +a large deployment the first backfill is minutes of work. """ from collections.abc import Awaitable, Callable @@ -27,7 +30,6 @@ if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient -_REPLAY_DAYS: Final = 1 GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend" # The unique constraint, in constraint order. NULL never matches itself in a unique index, so # every column is normalized to '' or the same group would be inserted again on every run. @@ -69,15 +71,26 @@ def _reconcile_day_sql() -> str: RECONCILE_DAY_SQL: Final = _reconcile_day_sql() +_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now" +_ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"' +# Pod clocks drift from the database clock and from each other, so rows are picked up from a +# little before the previous scan; rewriting a day twice is idempotent. _PENDING_DAYS_SQL: Final = ( - 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" >= $1 AND "date" <= $2 ORDER BY "date"' + 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ' + 'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') ' + 'ORDER BY "date"' ) class ReconciledThrough(BaseModel): + """``reconciled_through`` is the last closed UTC day the global table covers. ``scanned_at`` is + the database clock when the scan behind the last fully successful run started: every per-key + row written before it, on any day through the marker, is in the global table.""" + model_config = ConfigDict(frozen=True, extra="ignore") reconciled_through: str + scanned_at: str | None = None class _MarkerRow(BaseModel): @@ -92,6 +105,12 @@ class _DateRow(BaseModel): date: str +class _NowRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + now: str + + @dataclass(frozen=True, slots=True) class ReconcileResult: days_reconciled: tuple[str, ...] @@ -99,49 +118,70 @@ class ReconcileResult: failed_day: str | None = None -def _marker_from_param_value(value: object) -> str | None: +@dataclass(frozen=True, slots=True) +class _PendingScan: + marker: ReconciledThrough | None + scanned_at: str + days: tuple[str, ...] + + +def _marker_from_param_value(value: object) -> ReconciledThrough | None: try: - parsed: Final = ( + return ( ReconciledThrough.model_validate_json(value) if isinstance(value, str) else ReconciledThrough.model_validate(value) ) except ValidationError: return None - return parsed.reconciled_through -async def reconciled_through(prisma_client: "PrismaClient") -> str | None: - """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" +async def read_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: from litellm.proxy.utils import get_config_param row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value) -async def _record_reconciled_through(prisma_client: "PrismaClient", day: str) -> None: +async def reconciled_through(prisma_client: "PrismaClient") -> str | None: + """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" + marker: Final = await read_marker(prisma_client) + return None if marker is None else marker.reconciled_through + + +async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThrough) -> None: from litellm.proxy.utils import invalidate_config_param await ConfigRepository(prisma_client).set_param( - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ReconciledThrough(reconciled_through=day).model_dump_json() + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, marker.model_dump_json() ) await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -def _first_pending_day(marker: str | None) -> str: - if marker is None: - return "" - return (date.fromisoformat(marker) - timedelta(days=_REPLAY_DAYS)).isoformat() +async def _db_now(prisma_client: "PrismaClient") -> str: + rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) + return _NowRow.model_validate(rows[0]).now + + +async def _scan_pending(prisma_client: "PrismaClient", today: date) -> _PendingScan: + """Every closed UTC day (strictly before today) still to roll up, oldest first: days past the + marker, plus any day with per-key rows written since the scan behind the marker. Before a + run has fully succeeded there is no such scan, so every closed day is rolled up.""" + marker: Final = await read_marker(prisma_client) + scanned_at: Final = await _db_now(prisma_client) + last_closed_day: Final = (today - timedelta(days=1)).isoformat() + rows: Final = ( + await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day) + if marker is None or marker.scanned_at is None + else await prisma_client.db.query_raw( + _PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at + ) + ) + return _PendingScan(marker, scanned_at, tuple(_DateRow.model_validate(row).date for row in rows)) async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: - """Every closed UTC day (strictly before today) still to roll up, oldest first. The marker - day and the one before it are replayed so per-key rows that landed after their day was - rolled up (a flush straddling midnight, a late retry) are folded in.""" - marker: Final = await reconciled_through(prisma_client) - last_closed_day: Final = (today - timedelta(days=1)).isoformat() - rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), last_closed_day) - return tuple(_DateRow.model_validate(row).date for row in rows) + return (await _scan_pending(prisma_client, today)).days async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: @@ -155,26 +195,42 @@ async def run_daily_global_spend_reconcile( today: date | None = None, ) -> ReconcileResult: """Roll up every pending day, advancing the marker after each; a failing day stops the run - with the marker on the last good day so the next run resumes there.""" + with the marker on the last good day so the next run resumes there. The scan time is only + recorded once every pending day is done, so late rows a failed run saw are found again.""" effective_today: Final = today or datetime.now(timezone.utc).date() - days: Final = await pending_days(prisma_client, effective_today) - done: Final = await _reconcile_until_failure(prisma_client, days) - failed: Final = days[len(done)] if len(done) < len(days) else None - marker: Final = done[-1] if done else await reconciled_through(prisma_client) - return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=failed) + scan: Final = await _scan_pending(prisma_client, effective_today) + done: Final = await _reconcile_until_failure(prisma_client, scan) + if len(done) < len(scan.days): + marker: Final = await reconciled_through(prisma_client) + return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) + if scan.marker is not None or done: + await _record_marker(prisma_client, _advanced(scan.marker, done, scanned_at=scan.scanned_at)) + return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) -async def _reconcile_until_failure(prisma_client: "PrismaClient", days: tuple[str, ...]) -> tuple[str, ...]: - for index, day in enumerate(days): - if not await _reconcile_and_record(prisma_client, day): - return days[:index] - return days +def _advanced(marker: ReconciledThrough | None, days: tuple[str, ...], *, scanned_at: str | None) -> ReconciledThrough: + """The marker after ``days`` were rewritten: a late old day never moves it back.""" + through: Final = max((marker.reconciled_through if marker is not None else "", *days)) + return ReconciledThrough(reconciled_through=through, scanned_at=scanned_at) -async def _reconcile_and_record(prisma_client: "PrismaClient", day: str) -> bool: +async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: + for index, day in enumerate(scan.days): + if not await _reconcile_and_record(prisma_client, scan.marker, scan.days[: index + 1]): + return scan.days[:index] + return scan.days + + +async def _reconcile_and_record( + prisma_client: "PrismaClient", marker: ReconciledThrough | None, done_with_this: tuple[str, ...] +) -> bool: + day: Final = done_with_this[-1] try: await reconcile_day(prisma_client, day) - await _record_reconciled_through(prisma_client, day) + await _record_marker( + prisma_client, + _advanced(marker, done_with_this, scanned_at=None if marker is None else marker.scanned_at), + ) except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) return False diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 11ca72e7b3d..9a098744f08 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -15,6 +15,7 @@ from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( RECONCILE_DAY_SQL, + read_marker, reconciled_through, run_daily_global_spend_reconcile, run_scheduled_daily_global_spend_reconcile, @@ -41,13 +42,25 @@ class _FakeConfigTable: class _FakeDb: + """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, + so "rows written since the last scan" behaves like Postgres would.""" + def __init__(self, prisma: "_FakePrisma") -> None: self._prisma = prisma self.litellm_config = _FakeConfigTable() async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: - first, last = params - return [{"date": d} for d in sorted(self._prisma.user_days) if first <= d <= last] + if sql.startswith("SELECT (NOW()"): + self._prisma.clock += 1 + return [{"now": f"clock-{self._prisma.clock:04d}"}] + rows = self._prisma.user_rows + if len(params) == 1: + (last,) = params + return [{"date": d} for d in sorted(rows) if d <= last] + last, marker, scanned_at = params + return [ + {"date": d} for d, written in sorted(rows.items()) if d <= last and (d > marker or written >= scanned_at) + ] async def execute_raw(self, sql: str, *params: str) -> int: (day,) = params @@ -61,11 +74,17 @@ class _FakePrisma: """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: - self.user_days = user_days + self.clock = 0 + self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} self.failing_days = failing_days self.reconciled: list[str] = [] self.db = _FakeDb(self) + def write_late_row(self, day: str) -> None: + """A per-key row for ``day`` lands now, after whatever scans already happened.""" + self.clock += 1 + self.user_rows[day] = f"clock-{self.clock:04d}" + async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: stored = self.db.litellm_config.rows.get(value) return None if stored is None else _FakeConfigRow(value, stored) @@ -94,20 +113,66 @@ async def test_first_run_rolls_up_every_closed_day_and_never_today(): @pytest.mark.asyncio -async def test_later_run_replays_the_marker_day_and_the_day_before_only(): - """Days older than marker-1 are settled; the marker day and its predecessor are replayed so - per-key rows that landed after their day was rolled up get folded in.""" +async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.reconciled.clear() result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14") - assert "2026-09-01" not in prisma.reconciled + assert result.days_reconciled == ("2026-09-14",) assert await reconciled_through(prisma) == "2026-09-14" +@pytest.mark.asyncio +async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_run(): + """Per-key rows carry the request start date, so a delayed flush or retry can add spend to a + day far behind the marker. That day is rewritten, and the marker never moves back for it.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13")) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.reconciled.clear() + prisma.write_late_row("2026-09-01") + prisma.write_late_row("2026-09-03") + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-03") + assert "2026-09-05" not in prisma.reconciled + assert await reconciled_through(prisma) == "2026-09-13" + + +@pytest.mark.asyncio +async def test_a_late_row_seen_by_a_failed_run_is_seen_again_by_the_next_one(): + """The scan time only advances when every pending day was rewritten, otherwise a late row + found by the failed run would be counted as handled.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.write_late_row("2026-09-01") + prisma.failing_days = frozenset({"2026-09-01"}) + failed = await run_daily_global_spend_reconcile(prisma, today=TODAY) + prisma.failing_days = frozenset() + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert failed.failed_day == "2026-09-01" + assert failed.reconciled_through == "2026-09-13" + assert result.days_reconciled == ("2026-09-01",) + assert result.failed_day is None + + +@pytest.mark.asyncio +async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) + prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-13"}' + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-13") + marker = await read_marker(prisma) + assert marker is not None and marker.reconciled_through == "2026-09-13" and marker.scanned_at is not None + + @pytest.mark.asyncio async def test_a_run_with_no_new_closed_days_keeps_the_marker(): prisma = _FakePrisma(user_days=("2026-09-13",)) @@ -116,7 +181,7 @@ async def test_a_run_with_no_new_closed_days_keeps_the_marker(): result = await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) - assert result.days_reconciled == ("2026-09-13",) + assert result.days_reconciled == () assert result.reconciled_through == "2026-09-13" @@ -149,11 +214,10 @@ async def test_the_next_run_resumes_from_the_failed_day(): @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): - """A late flush for the day before the marker is exactly the replay case; when that replay - fails the marker must stay put and the operator must hear about it.""" + """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" prisma = _FakePrisma(user_days=("2026-09-13",)) await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) - prisma.user_days = ("2026-09-12", "2026-09-13") + prisma.write_late_row("2026-09-12") prisma.failing_days = frozenset({"2026-09-12"}) alert = AsyncMock() From 5b5bbac769e548199393e54aacf346953c5c5528 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:16:21 +0000 Subject: [PATCH 021/102] fix(team): link new members to the shared team member budget so /team/update applies to them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_helpers/utils.py | 23 +-- .../test_management_helpers_utils.py | 135 +++++++++--------- 2 files changed, 82 insertions(+), 76 deletions(-) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index f3bd4b0f6dd..2e7458232cc 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -291,9 +291,9 @@ async def _clone_team_default_budget_for_member( member budget. Returns the new budget_id, or None if the default budget no longer exists in the DB. - Used when adding a new team member without an explicit per-member budget, - so the member starts with the team default's values but gets their own - private budget row (which can be edited independently). + Used when adding a new team member with a per-member ``budget_duration`` + but no other per-member limit, so the member keeps the team default's + values in their own private budget row while the reset window differs. ``budget_duration_override`` replaces the default's reset window for this member while keeping the default's other limits, so an admin can set a @@ -344,14 +344,21 @@ async def _resolve_member_budget_id( """ Resolve the budget a new team member should be linked to. - Explicit per-member limits create a fresh budget. Otherwise the team's - default member budget is cloned (with ``budget_duration`` overriding its - reset window while keeping its other limits). A lone ``budget_duration`` - with no team default creates a window-only budget. With nothing set the - member gets no budget. + Explicit per-member limits create a fresh budget. Otherwise the member is + linked to the team's shared default member budget, so later ``/team/update`` + changes reach them; ``/team/member_update`` clones that row on first write. + A lone ``budget_duration`` clones the default with the reset window + overridden, or creates a window-only budget when there is no team default. + With nothing set the member gets no budget. """ has_explicit_limit: Final = max_budget_in_team is not None or allowed_models is not None + if not has_explicit_limit and default_team_budget_id is not None and budget_duration is None: + default_budget: Final = await _budget_table(prisma_client, tx).find_unique( + where={"budget_id": default_team_budget_id} + ) + return default_team_budget_id if default_budget is not None else None + if not has_explicit_limit and default_team_budget_id is not None: return await _clone_team_default_budget_for_member( prisma_client=prisma_client, diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index a6b1fc32eda..00de5171aa7 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -164,13 +164,16 @@ async def test_management_otel_span_redacts_nested_submission_env_var_secrets( @pytest.mark.asyncio -async def test_add_new_member_clones_default_team_budget_id(): +async def test_add_new_member_links_default_team_budget_id(): """ - Test that add_new_member CLONES the team's default member budget when - max_budget_in_team is None and a default_team_budget_id is provided. + A member added without any per-member limit must be LINKED to the team's + shared default member budget, not given a private copy of it. - Cloning (rather than sharing the same budget row) is what lets admins later - edit one member's budget without mutating every other member's budget. + Linking is what makes a later ``/team/update team_member_budget=...`` + reach existing members: the auth check reads the budget row behind the + membership, so a private clone would freeze the member at the old cap. + Per-member isolation is handled by ``/team/member_update`` cloning the + shared row on first write. """ from litellm.proxy._types import LitellmUserRoles @@ -178,7 +181,6 @@ async def test_add_new_member_clones_default_team_budget_id(): test_user_id = "test_user_123" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" - test_cloned_budget_id = "cloned_budget_xyz" test_admin_name = "test_admin" new_member = Member(user_id=test_user_id, role="user") @@ -202,36 +204,19 @@ async def test_add_new_member_clones_default_team_budget_id(): return_value=mock_user_response ) - # Mock the default budget row fetched for cloning. mock_default_budget_row = MagicMock() - mock_default_budget_row.model_dump.return_value = { - "budget_id": test_default_budget_id, - "max_budget": 100.0, - "soft_budget": None, - "max_parallel_requests": None, - "tpm_limit": 1000, - "rpm_limit": None, - "model_max_budget": None, - "budget_duration": "1d", - "allowed_models": [], - } + mock_default_budget_row.budget_id = test_default_budget_id mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=mock_default_budget_row ) - - # Mock the cloned budget row that .create() returns. - mock_cloned_budget_row = MagicMock() - mock_cloned_budget_row.budget_id = test_cloned_budget_id - mock_prisma_client.db.litellm_budgettable.create = AsyncMock( - return_value=mock_cloned_budget_row - ) + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() # Mock the team membership creation mock_team_membership_response = MagicMock() mock_team_membership_response.model_dump.return_value = { "team_id": test_team_id, "user_id": test_user_id, - "budget_id": test_cloned_budget_id, + "budget_id": test_default_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( @@ -251,33 +236,67 @@ async def test_add_new_member_clones_default_team_budget_id(): assert result_user is not None assert result_user.user_id == test_user_id - # Membership should be linked to the new cloned budget, not the shared default. + # Membership points at the shared default row itself. assert result_team_membership is not None - assert result_team_membership.budget_id == test_cloned_budget_id - assert result_team_membership.budget_id != test_default_budget_id + assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_teammembership.create.assert_called_once() - # The clone must have happened: find_unique on the default, create for the clone. + # The default is only checked for existence; no private budget row is created. mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": test_default_budget_id} ) - mock_prisma_client.db.litellm_budgettable.create.assert_called_once() - cloned_create_data = ( - mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs["data"] - ) - # Cloned values from the default budget row - assert cloned_create_data["max_budget"] == 100.0 - assert cloned_create_data["tpm_limit"] == 1000 - assert cloned_create_data["budget_duration"] == "1d" - assert cloned_create_data["created_by"] == user_api_key_dict.user_id + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() team_membership_call_args = ( mock_prisma_client.db.litellm_teammembership.create.call_args ) create_data = team_membership_call_args.kwargs["data"] - assert create_data["budget_id"] == test_cloned_budget_id + assert create_data["budget_id"] == test_default_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_no_budget_when_default_budget_row_is_missing(): + """If team metadata still names a default member budget whose row was + deleted, the member must get no budget rather than a dangling link that + the membership foreign key would reject.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="missing-default-user", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "missing-default-user", + "user_email": None, + "teams": ["team-md"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mock_user_response + ) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() + mock_prisma_client.db.litellm_teammembership.create = AsyncMock() + + _, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-md", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id="deleted-default", + ) + + assert result_team_membership is None + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() + mock_prisma_client.db.litellm_teammembership.create.assert_not_called() @pytest.mark.asyncio @@ -636,18 +655,17 @@ async def test_add_new_member_persists_budget_duration_without_max_budget(): @pytest.mark.asyncio -async def test_add_new_member_with_user_email_clones_default_budget(): +async def test_add_new_member_with_user_email_links_default_budget(): """ Test add_new_member with user_email instead of user_id and a team default - budget. The default budget should be CLONED into a new private row for - this user, not shared with other members of the team. + budget. The membership must link the shared default row so team-level + budget updates keep applying to this member. """ from litellm.proxy._types import LitellmUserRoles test_user_email = "test@example.com" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" - test_cloned_budget_id = "cloned_budget_for_email_user" test_admin_name = "test_admin" new_member = Member(user_email=test_user_email, role="user") @@ -669,35 +687,18 @@ async def test_add_new_member_with_user_email_clones_default_budget(): } mock_prisma_client.insert_data = AsyncMock(return_value=mock_user_response) - # Default budget that will be cloned mock_default_budget_row = MagicMock() - mock_default_budget_row.model_dump.return_value = { - "budget_id": test_default_budget_id, - "max_budget": 25.0, - "soft_budget": None, - "max_parallel_requests": None, - "tpm_limit": None, - "rpm_limit": None, - "model_max_budget": None, - "budget_duration": None, - "allowed_models": [], - } + mock_default_budget_row.budget_id = test_default_budget_id mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=mock_default_budget_row ) - - # Cloned budget result - mock_cloned_budget_row = MagicMock() - mock_cloned_budget_row.budget_id = test_cloned_budget_id - mock_prisma_client.db.litellm_budgettable.create = AsyncMock( - return_value=mock_cloned_budget_row - ) + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() mock_team_membership_response = MagicMock() mock_team_membership_response.model_dump.return_value = { "team_id": test_team_id, "user_id": "generated_user_id", - "budget_id": test_cloned_budget_id, + "budget_id": test_default_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( @@ -717,9 +718,8 @@ async def test_add_new_member_with_user_email_clones_default_budget(): assert result_user is not None assert result_user.user_email == test_user_email - # Membership should point at the cloned (private) budget, not the shared default. assert result_team_membership is not None - assert result_team_membership.budget_id == test_cloned_budget_id + assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.get_data.assert_called_once_with( key_val={"user_email": test_user_email}, @@ -733,11 +733,10 @@ async def test_add_new_member_with_user_email_clones_default_budget(): assert insert_data["user_email"] == test_user_email assert insert_data["teams"] == [test_team_id] - # Confirm the clone path ran mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": test_default_budget_id} ) - mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() @pytest.mark.asyncio From ab99be9dad023d7d5e25d3e9352f7954357d4963 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:43:21 +0000 Subject: [PATCH 022/102] test(team): cover team_member_budget propagation and per-member isolation end to end Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_management_helpers_utils.py | 172 +++++++++++++++--- 1 file changed, 148 insertions(+), 24 deletions(-) diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index 00de5171aa7..7512a3dfae8 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -1,13 +1,17 @@ import json +from collections.abc import Mapping from datetime import datetime, timezone -from litellm._uuid import uuid -from unittest.mock import AsyncMock, MagicMock +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest - +import litellm +from litellm._uuid import uuid from litellm.proxy._types import ( + LiteLLM_BudgetTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, LiteLLM_UserTable, Member, UserAPIKeyAuth, @@ -165,19 +169,8 @@ async def test_management_otel_span_redacts_nested_submission_env_var_secrets( @pytest.mark.asyncio async def test_add_new_member_links_default_team_budget_id(): - """ - A member added without any per-member limit must be LINKED to the team's - shared default member budget, not given a private copy of it. - - Linking is what makes a later ``/team/update team_member_budget=...`` - reach existing members: the auth check reads the budget row behind the - membership, so a private clone would freeze the member at the old cap. - Per-member isolation is handled by ``/team/member_update`` cloning the - shared row on first write. - """ from litellm.proxy._types import LitellmUserRoles - # Setup test data test_user_id = "test_user_123" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" @@ -236,14 +229,12 @@ async def test_add_new_member_links_default_team_budget_id(): assert result_user is not None assert result_user.user_id == test_user_id - # Membership points at the shared default row itself. assert result_team_membership is not None assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_teammembership.create.assert_called_once() - # The default is only checked for existence; no private budget row is created. mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": test_default_budget_id} ) @@ -258,9 +249,6 @@ async def test_add_new_member_links_default_team_budget_id(): @pytest.mark.asyncio async def test_add_new_member_no_budget_when_default_budget_row_is_missing(): - """If team metadata still names a default member budget whose row was - deleted, the member must get no budget rather than a dangling link that - the membership foreign key would reject.""" from litellm.proxy._types import LitellmUserRoles new_member = Member(user_id="missing-default-user", role="user") @@ -656,11 +644,6 @@ async def test_add_new_member_persists_budget_duration_without_max_budget(): @pytest.mark.asyncio async def test_add_new_member_with_user_email_links_default_budget(): - """ - Test add_new_member with user_email instead of user_id and a team default - budget. The membership must link the shared default row so team-level - budget updates keep applying to this member. - """ from litellm.proxy._types import LitellmUserRoles test_user_email = "test@example.com" @@ -739,6 +722,147 @@ async def test_add_new_member_with_user_email_links_default_budget(): mock_prisma_client.db.litellm_budgettable.create.assert_not_called() +class _FakeBudgetTable: + def __init__(self) -> None: + self.rows: dict[str, dict[str, object]] = {} + + def _record(self, budget_id: str) -> LiteLLM_BudgetTable: + row: Final = self.rows[budget_id] + return LiteLLM_BudgetTable(**{k: v for k, v in row.items() if k in LiteLLM_BudgetTable.model_fields}) + + async def create( + self, *, data: Mapping[str, object], include: Mapping[str, bool] | None = None + ) -> LiteLLM_BudgetTable: + budget_id: Final = str(data.get("budget_id") or uuid.uuid4()) + self.rows[budget_id] = {**data, "budget_id": budget_id} + return self._record(budget_id) + + async def find_unique(self, *, where: Mapping[str, str]) -> LiteLLM_BudgetTable | None: + return self._record(where["budget_id"]) if where["budget_id"] in self.rows else None + + async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> LiteLLM_BudgetTable: + self.rows[where["budget_id"]] = {**self.rows[where["budget_id"]], **data} + return self._record(where["budget_id"]) + + +class _FakeMembershipTable: + def __init__(self, budgets: _FakeBudgetTable) -> None: + self.budgets: Final = budgets + self.budget_ids: dict[tuple[str, str], str | None] = {} + + def membership(self, team_id: str, user_id: str) -> LiteLLM_TeamMembership: + budget_id: Final = self.budget_ids[(team_id, user_id)] + return LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + budget_id=budget_id, + litellm_budget_table=self.budgets._record(budget_id) if budget_id is not None else None, + ) + + async def create(self, *, data: Mapping[str, str], include: Mapping[str, bool]) -> LiteLLM_TeamMembership: + self.budget_ids[(data["team_id"], data["user_id"])] = data["budget_id"] + return self.membership(data["team_id"], data["user_id"]) + + async def upsert(self, *, where: Mapping[str, Mapping[str, str]], data: Mapping[str, Mapping[str, object]]) -> None: + key: Final = where["user_id_team_id"] + connect: Final = data["update"]["litellm_budget_table"] + assert isinstance(connect, dict) + self.budget_ids[(key["team_id"], key["user_id"])] = connect["connect"]["budget_id"] + + +class _FakeUserTable: + async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id=where["user_id"], teams=list(data["create"].get("teams", []))) + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: + return 1 + + +class _FakeDb: + def __init__(self) -> None: + self.litellm_budgettable: Final = _FakeBudgetTable() + self.litellm_teammembership: Final = _FakeMembershipTable(self.litellm_budgettable) + self.litellm_usertable: Final = _FakeUserTable() + + +@pytest.mark.asyncio +async def test_team_update_reaches_inherited_members_but_not_overridden_ones(): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.auth_checks import _check_team_member_budget + from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership + from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.utils import ProxyLogging + + db: Final = _FakeDb() + prisma_client: Final = MagicMock() + prisma_client.db = db + admin: Final = UserAPIKeyAuth(user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN) + team_id: Final = "team-shared-default" + default_budget: Final = await db.litellm_budgettable.create(data={"budget_id": "team-default", "max_budget": 100.0}) + team: Final = LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": default_budget.budget_id}) + + for user_id in ("inherits", "overridden"): + await add_new_member( + new_member=Member(user_id=user_id, role="user"), + max_budget_in_team=None, + prisma_client=prisma_client, + team_id=team_id, + user_api_key_dict=admin, + litellm_proxy_admin_name="admin", + default_team_budget_id=default_budget.budget_id, + ) + + await _upsert_budget_and_membership( + db, + team_id=team_id, + user_id="overridden", + existing_budget_id=default_budget.budget_id, + user_api_key_dict=admin, + budget_patch={"max_budget": 50.0}, + team_default_budget_id=default_budget.budget_id, + ) + assert db.litellm_teammembership.membership(team_id, "inherits").budget_id == default_budget.budget_id + assert db.litellm_teammembership.membership(team_id, "overridden").budget_id != default_budget.budget_id + assert db.litellm_budgettable.rows[default_budget.budget_id]["max_budget"] == 100.0 + + with patch( # test-quality-ok: update_budget reads this module global; no dependency injection seam exists + "litellm.proxy.proxy_server.prisma_client", prisma_client + ): + await TeamMemberBudgetHandler.upsert_team_member_budget_table( + team_table=team, + user_api_key_dict=admin, + updated_kv={}, + team_member_budget=1.0, + ) + + async def spend_from_membership(counter_key: str, fallback_spend: float, max_budget: float | None = None) -> float: + return fallback_spend + + async def check(user_id: str, spend: float) -> None: + membership: Final = db.litellm_teammembership.membership(team_id, user_id).model_copy(update={"spend": spend}) + with patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists + "litellm.proxy.proxy_server.get_current_spend", spend_from_membership + ): + await _check_team_member_budget( + team_object=team, + user_object=LiteLLM_UserTable(user_id=user_id), + valid_token=UserAPIKeyAuth(token="tok", user_id=user_id, team_id=team_id), + prisma_client=prisma_client, + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + team_membership=membership, + team_membership_loaded=True, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await check("inherits", spend=2.0) + assert exc_info.value.max_budget == 1.0 + await check("overridden", spend=2.0) + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await check("overridden", spend=60.0) + assert exc_info.value.max_budget == 50.0 + + @pytest.mark.asyncio async def test_attach_object_permission_to_dict_with_object_permission_id(): """ From 0601d2bb03646c596a680d580b0f9bb5a83ee237 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 01:58:37 +0000 Subject: [PATCH 023/102] feat(proxy): carry response time metrics through LiteLLM_DailyGlobalSpend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 2 ++ .../litellm_proxy_extras/schema.prisma | 2 ++ .../management_endpoints/common_daily_activity.py | 2 ++ litellm/proxy/schema.prisma | 2 ++ .../spend_tracking/daily_global_spend_rollup.py | 2 ++ schema.prisma | 2 ++ .../test_common_daily_activity.py | 6 ++++++ .../test_daily_global_spend_rollup.py | 13 +++++++++++-- 8 files changed, 29 insertions(+), 2 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql index 1d6cdea0c7b..d0bc3e159de 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql @@ -20,6 +20,8 @@ CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGlobalSpend" ( "api_requests" BIGINT NOT NULL DEFAULT 0, "successful_requests" BIGINT NOT NULL DEFAULT 0, "failed_requests" BIGINT NOT NULL DEFAULT 0, + "total_response_time_ms" BIGINT NOT NULL DEFAULT 0, + "timed_requests" BIGINT NOT NULL DEFAULT 0, "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, "updated_at" TIMESTAMP(3) NOT NULL, diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a73e8774c87..42769c323a9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -837,6 +837,8 @@ model LiteLLM_DailyGlobalSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 31ec0c51b7d..47465324f42 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -767,6 +767,8 @@ _KEY_FREE_SOURCE_COLUMNS: Final = ( "api_requests", "successful_requests", "failed_requests", + "total_response_time_ms", + "timed_requests", ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a73e8774c87..42769c323a9 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -837,6 +837,8 @@ model LiteLLM_DailyGlobalSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index 73068381dab..ba4a4e4e3d6 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -43,6 +43,8 @@ _METRIC_COLUMNS: Final = ( "api_requests", "successful_requests", "failed_requests", + "total_response_time_ms", + "timed_requests", "compression_savings_spend", "prompt_caching_savings_spend", "gateway_injected_caching_savings_spend", diff --git a/schema.prisma b/schema.prisma index a73e8774c87..42769c323a9 100644 --- a/schema.prisma +++ b/schema.prisma @@ -837,6 +837,8 @@ model LiteLLM_DailyGlobalSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 2f25c507e0f..ee9f886acec 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1771,6 +1771,10 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_ ] _seed_daily_user_spend(_aggregated_postgresql, rows) with _aggregated_postgresql.cursor() as cur: + cur.execute( + 'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms = prompt_tokens * 25, ' + "timed_requests = api_requests" + ) cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal cur.execute( re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg @@ -1805,6 +1809,8 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_ assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 3 assert from_global.model_dump() == from_per_key.model_dump() assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + assert from_global.metadata.total_response_time_ms == 2 * n_keys * 10 * 25 + assert from_global.metadata.total_timed_requests == 2 * n_keys assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 9a098744f08..b5b78229c11 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -336,6 +336,8 @@ _DAILY_USER_SPEND_DDL: Final = """ api_requests BIGINT DEFAULT 0, successful_requests BIGINT DEFAULT 0, failed_requests BIGINT DEFAULT 0, + total_response_time_ms BIGINT DEFAULT 0, + timed_requests BIGINT DEFAULT 0, created_at TIMESTAMP DEFAULT now(), updated_at TIMESTAMP, UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) @@ -345,12 +347,14 @@ _DAILY_USER_SPEND_DDL: Final = """ _PER_KEY_SUMS_SQL: Final = """ SELECT COALESCE(model, '') AS model, COALESCE(model_group, '') AS model_group, COALESCE(custom_llm_provider, '') AS custom_llm_provider, - SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests + SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests, + SUM(total_response_time_ms) AS total_response_time_ms, SUM(timed_requests) AS timed_requests FROM "LiteLLM_DailyUserSpend" WHERE date = %s GROUP BY 1, 2, 3 ORDER BY 1, 2, 3 """ _GLOBAL_ROWS_SQL: Final = """ - SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests + SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests, + total_response_time_ms, timed_requests FROM "LiteLLM_DailyGlobalSpend" WHERE date = %s ORDER BY 1, 2, 3 """ @@ -380,6 +384,8 @@ def _user_txn(**overrides): "api_requests": 1, "successful_requests": 1, "failed_requests": 0, + "total_response_time_ms": 800, + "timed_requests": 1, **overrides, } @@ -393,6 +399,8 @@ def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: float(r["spend"]), int(r["prompt_tokens"]), int(r["api_requests"]), + int(r["total_response_time_ms"]), + int(r["timed_requests"]), ) # pyright: ignore[reportArgumentType] # dict_row values are untyped for r in rows ] @@ -436,5 +444,6 @@ def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_p assert _normalized(global_rows) == _normalized(per_key) assert sum(float(r["spend"]) for r in global_rows) == pytest.approx(15.0) # pyright: ignore[reportArgumentType] # dict_row values are untyped + assert sum(int(r["total_response_time_ms"]) for r in global_rows) == 1600 # pyright: ignore[reportArgumentType] # dict_row values are untyped assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")] assert untouched == [] From 834313af4b188ee5561e8c0e8094fad399e5ac78 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 12:13:57 +0000 Subject: [PATCH 024/102] test(proxy): assert the global rollup split and scheduler through behavior, not SQL text or add_job arguments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_common_daily_activity.py | 95 ++++++++++--------- .../proxy/proxy_server/test_lifecycle.py | 24 +++-- 2 files changed, 66 insertions(+), 53 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index ac761315a27..fc3ede88aa9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1647,30 +1647,6 @@ async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -def test_aggregated_sql_splits_the_key_free_arm_at_the_marker_and_keeps_the_key_arm_per_key(): - sql, params = _build_aggregated_sql_query(**_unfiltered_user_query(), global_rollup_through="2026-06-01") - marker_param: Final = f"${len(params)}" - - assert params[-1] == "2026-06-01" - assert ( - f'FROM "LiteLLM_DailyGlobalSpend"\n WHERE date >= $1 AND date <= $2 AND date <= {marker_param}' - in sql - ) - assert ( - f'FROM "LiteLLM_DailyUserSpend"\n WHERE date >= $1 AND date <= $2 AND date > {marker_param}' in sql - ) - key_arm: Final = sql.split("UNION ALL\n (WITH top_api_keys")[1] - assert "LiteLLM_DailyGlobalSpend" not in key_arm - assert marker_param not in key_arm - - -def test_aggregated_sql_without_a_marker_reads_the_per_key_table_only(): - sql, params = _build_aggregated_sql_query(**_unfiltered_user_query()) - - assert "LiteLLM_DailyGlobalSpend" not in sql - assert params[-1] == PTU_SENTINEL_API_KEY - - _GLOBAL_SPEND_MIGRATION: Final = ( pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" @@ -1687,7 +1663,10 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_ ): """Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must give the same response as reading everything per-key: day 1 from the global table, day 2 - live, one grand total across both. The per-key arm stays on the user table throughout.""" + live, one grand total across both. Per-key rows that land after the rollup then tell the + two sources apart: a late day 1 row is invisible to totals until the next reconcile while a + late day 2 row shows up at once, and both keys rank in the key breakdown, which stays + per-key throughout.""" n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 rows: Final = [ ( @@ -1720,39 +1699,56 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_ ) _aggregated_postgresql.commit() - async def read(marker: str | None, sql_seen: list[str]): + async def read(marker: str | None): await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) prisma = _prisma_with_marker(marker) - run_query = _psycopg_query_raw(_aggregated_postgresql, []) - - async def query_raw(sql: str, *params: str): - sql_seen.append(sql) - return await run_query(sql, *params) - - prisma.db.query_raw = query_raw + prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, []) return await get_daily_activity_aggregated( prisma_client=prisma, entity_metadata_field=None, **_unfiltered_user_query(), ) - per_key_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim - global_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim - from_per_key = await read(None, per_key_sql) - from_global = await read("2026-06-01", global_sql) - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + from_per_key = await read(None) + from_global = await read("2026-06-01") - assert per_key_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 0 - assert global_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 1 - assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 3 assert from_global.model_dump() == from_per_key.model_dump() - assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + seeded_spend: Final = 2 * sum(float(i + 1) for i in range(n_keys)) + assert from_global.metadata.total_spend == pytest.approx(seeded_spend) assert from_global.metadata.total_response_time_ms == 2 * n_keys * 10 * 25 assert from_global.metadata.total_timed_requests == 2 * n_keys assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} + with _aggregated_postgresql.cursor() as cur: + cur.executemany( + """ + INSERT INTO "LiteLLM_DailyUserSpend" + (id, user_id, date, api_key, model, model_group, custom_llm_provider, + endpoint, prompt_tokens, spend, api_requests, successful_requests) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + """, + [ + ("late-1", "user-late", "2026-06-01", "key-late-1", "gpt-5", "", "openai", None, 10, 1000.0, 1, 1), + ("late-2", "user-late", "2026-06-02", "key-late-2", "gpt-5", "", "openai", None, 10, 500.0, 1, 1), + ], + ) + _aggregated_postgresql.commit() + + late_per_key = await read(None) + late_global = await read("2026-06-01") + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + assert late_per_key.metadata.total_spend == pytest.approx(seeded_spend + 1000.0 + 500.0) + assert late_global.metadata.total_spend == pytest.approx(seeded_spend + 500.0) + by_day: Final = {day.date.isoformat(): day for day in late_global.results} + assert by_day["2026-06-01"].metrics.spend == pytest.approx(seeded_spend / 2) + assert by_day["2026-06-02"].metrics.spend == pytest.approx(seeded_spend / 2 + 500.0) + assert by_day["2026-06-01"].breakdown.api_keys["key-late-1"].metrics.spend == pytest.approx(1000.0) + assert by_day["2026-06-02"].breakdown.api_keys["key-late-2"].metrics.spend == pytest.approx(500.0) + assert late_global.metadata.total_api_keys == n_keys + 2 + @pytest.mark.asyncio async def test_get_daily_activity_aggregated_reports_exact_limit_key_count_as_complete( @@ -1810,7 +1806,20 @@ async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_mo """Rows stored with an empty or NULL model_group must land in the model_groups breakdown under their model name instead of vanishing from the usage UI.""" rows: Final = [ - ("row-0", "user-0", "2026-06-01", "key-0", "gpt-5", "gpt-5-eu", "openai", "/v1/chat/completions", 10, 7.0, 1, 1), + ( + "row-0", + "user-0", + "2026-06-01", + "key-0", + "gpt-5", + "gpt-5-eu", + "openai", + "/v1/chat/completions", + 10, + 7.0, + 1, + 1, + ), ("row-1", "user-1", "2026-06-01", "key-1", "gpt-5", "", "openai", "/v1/chat/completions", 10, 3.0, 1, 1), ("row-2", "user-2", "2026-06-01", "key-2", "claude-x", None, "anthropic", "/v1/messages", 10, 2.0, 1, 1), ] diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index ee72e98ffa9..6121608b658 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -28,6 +28,7 @@ from typing import List, Optional, Union from unittest.mock import AsyncMock, MagicMock, patch import pytest +from apscheduler.schedulers.asyncio import AsyncIOScheduler from fastapi import FastAPI from pydantic import BaseModel from typing_extensions import TypedDict @@ -1042,8 +1043,8 @@ async def test_spend_report_locks_are_never_released(): proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited() -def _init_daily_global_spend_reconcile_job() -> tuple[MagicMock, MagicMock, MagicMock]: - scheduler = MagicMock() +def _init_daily_global_spend_reconcile_job() -> tuple[AsyncIOScheduler, MagicMock, MagicMock]: + scheduler = AsyncIOScheduler() proxy_logging_obj = MagicMock() proxy_logging_obj.alerting_handler = AsyncMock() prisma_client = MagicMock() @@ -1058,28 +1059,31 @@ def _init_daily_global_spend_reconcile_job() -> tuple[MagicMock, MagicMock, Magi def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run(): """Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a fresh deploy switches usage reads to the global table without waiting for the nightly - run, and replaces any previous registration of the same job id.""" + run, and after that it fires once a day at 00:30 UTC, when the previous UTC day is closed.""" from datetime import datetime, timedelta, timezone from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID scheduler, _, _ = _init_daily_global_spend_reconcile_job() + job = scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) + assert job is not None - (call,) = scheduler.add_job.call_args_list - assert call.kwargs["id"] == DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID - assert call.kwargs["replace_existing"] is True - assert call.args[1:] == ("cron",) - assert (call.kwargs["hour"], call.kwargs["minute"], call.kwargs["timezone"]) == (0, 30, "UTC") - assert timedelta(0) < call.kwargs["next_run_time"] - datetime.now(timezone.utc) <= timedelta(minutes=2) + assert timedelta(0) < job.next_run_time - datetime.now(timezone.utc) <= timedelta(minutes=2) + after_catch_up = datetime(2026, 9, 16, 12, 0, tzinfo=timezone.utc) + assert job.trigger.get_next_fire_time(None, after_catch_up) == datetime(2026, 9, 17, 0, 30, tzinfo=timezone.utc) + just_after_a_run = datetime(2026, 9, 17, 0, 30, 1, tzinfo=timezone.utc) + assert job.trigger.get_next_fire_time(None, just_after_a_run) == datetime(2026, 9, 18, 0, 30, tzinfo=timezone.utc) @pytest.mark.asyncio async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch): + from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID + scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job() run = AsyncMock() monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run) - await scheduler.add_job.call_args.args[0]() + await scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID).func() run.assert_awaited_once() assert run.await_args.args == (prisma_client,) From f25d65940d2a013d1604c729a7dc39836df5e31f Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 12:39:48 +0000 Subject: [PATCH 025/102] fix(proxy): take the closed-day cutoff for the global spend rollup from the database clock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../daily_global_spend_rollup.py | 43 +++++------- .../test_daily_global_spend_rollup.py | 70 +++++++++++-------- 2 files changed, 59 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index ba4a4e4e3d6..c135c7d1d9c 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -12,7 +12,7 @@ a large deployment the first backfill is minutes of work. from collections.abc import Awaitable, Callable from dataclasses import dataclass -from datetime import date, datetime, timedelta, timezone +from datetime import date, timedelta from typing import TYPE_CHECKING, Final from pydantic import BaseModel, ConfigDict, ValidationError @@ -73,7 +73,7 @@ def _reconcile_day_sql() -> str: RECONCILE_DAY_SQL: Final = _reconcile_day_sql() -_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now" +_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now, (NOW() AT TIME ZONE 'UTC')::date::text AS today" _ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"' # Pod clocks drift from the database clock and from each other, so rows are picked up from a # little before the previous scan; rewriting a day twice is idempotent. @@ -111,6 +111,7 @@ class _NowRow(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") now: str + today: str @dataclass(frozen=True, slots=True) @@ -160,18 +161,18 @@ async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThroug await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -async def _db_now(prisma_client: "PrismaClient") -> str: +async def _db_now(prisma_client: "PrismaClient") -> _NowRow: rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) - return _NowRow.model_validate(rows[0]).now + return _NowRow.model_validate(rows[0]) -async def _scan_pending(prisma_client: "PrismaClient", today: date) -> _PendingScan: - """Every closed UTC day (strictly before today) still to roll up, oldest first: days past the - marker, plus any day with per-key rows written since the scan behind the marker. Before a - run has fully succeeded there is no such scan, so every closed day is rolled up.""" +async def _scan_pending(prisma_client: "PrismaClient") -> _PendingScan: + """Every closed UTC day (strictly before the database's today) still to roll up, oldest first: + days past the marker, plus any day with per-key rows written since the scan behind the marker. + Before a run has fully succeeded there is no such scan, so every closed day is rolled up.""" marker: Final = await read_marker(prisma_client) - scanned_at: Final = await _db_now(prisma_client) - last_closed_day: Final = (today - timedelta(days=1)).isoformat() + db_now: Final = await _db_now(prisma_client) + last_closed_day: Final = (date.fromisoformat(db_now.today) - timedelta(days=1)).isoformat() rows: Final = ( await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day) if marker is None or marker.scanned_at is None @@ -179,11 +180,11 @@ async def _scan_pending(prisma_client: "PrismaClient", today: date) -> _PendingS _PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at ) ) - return _PendingScan(marker, scanned_at, tuple(_DateRow.model_validate(row).date for row in rows)) + return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows)) -async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: - return (await _scan_pending(prisma_client, today)).days +async def pending_days(prisma_client: "PrismaClient") -> tuple[str, ...]: + return (await _scan_pending(prisma_client)).days async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: @@ -192,15 +193,11 @@ async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day) -async def run_daily_global_spend_reconcile( - prisma_client: "PrismaClient", - today: date | None = None, -) -> ReconcileResult: +async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> ReconcileResult: """Roll up every pending day, advancing the marker after each; a failing day stops the run with the marker on the last good day so the next run resumes there. The scan time is only recorded once every pending day is done, so late rows a failed run saw are found again.""" - effective_today: Final = today or datetime.now(timezone.utc).date() - scan: Final = await _scan_pending(prisma_client, effective_today) + scan: Final = await _scan_pending(prisma_client) done: Final = await _reconcile_until_failure(prisma_client, scan) if len(done) < len(scan.days): marker: Final = await reconciled_through(prisma_client) @@ -243,13 +240,12 @@ async def run_scheduled_daily_global_spend_reconcile( prisma_client: "PrismaClient", pod_lock_manager: "PodLockManager | None" = None, alert: Callable[[str], Awaitable[None]] | None = None, - today: date | None = None, ) -> ReconcileResult | None: """Run the reconcile under a cross-pod lock so one proxy does the work; the lock only saves effort (each day is an idempotent rewrite), so an unreachable Redis runs unguarded rather than skipping.""" redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache if pod_lock_manager is None or redis_cache is None: - return await _run_and_alert(prisma_client, alert=alert, today=today) + return await _run_and_alert(prisma_client, alert=alert) acquired: Final = await pod_lock_manager.acquire_lock( cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, ttl=DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS @@ -258,7 +254,7 @@ async def run_scheduled_daily_global_spend_reconcile( verbose_proxy_logger.info("Daily global spend reconcile: another pod holds the lock, skipping this run") return None try: - return await _run_and_alert(prisma_client, alert=alert, today=today) + return await _run_and_alert(prisma_client, alert=alert) finally: if acquired: await pod_lock_manager.release_lock(cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) @@ -277,9 +273,8 @@ async def _run_and_alert( prisma_client: "PrismaClient", *, alert: Callable[[str], Awaitable[None]] | None, - today: date | None, ) -> ReconcileResult: - result: Final = await run_daily_global_spend_reconcile(prisma_client, today=today) + result: Final = await run_daily_global_spend_reconcile(prisma_client) if result.days_reconciled: verbose_proxy_logger.info( "Daily global spend reconcile: rolled up %d day(s), reconciled through %s", diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index b5b78229c11..9655953134a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -43,7 +43,8 @@ class _FakeConfigTable: class _FakeDb: """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, - so "rows written since the last scan" behaves like Postgres would.""" + so "rows written since the last scan" behaves like Postgres would. The database's own + date decides which day is still open, never the pod's clock.""" def __init__(self, prisma: "_FakePrisma") -> None: self._prisma = prisma @@ -52,7 +53,7 @@ class _FakeDb: async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: if sql.startswith("SELECT (NOW()"): self._prisma.clock += 1 - return [{"now": f"clock-{self._prisma.clock:04d}"}] + return [{"now": f"clock-{self._prisma.clock:04d}", "today": self._prisma.today.isoformat()}] rows = self._prisma.user_rows if len(params) == 1: (last,) = params @@ -73,8 +74,11 @@ class _FakeDb: class _FakePrisma: """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" - def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: + def __init__( + self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset(), today: date = TODAY + ) -> None: self.clock = 0 + self.today = today self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} self.failing_days = failing_days self.reconciled: list[str] = [] @@ -98,12 +102,14 @@ async def _fresh_marker_cache(): @pytest.mark.asyncio -async def test_first_run_rolls_up_every_closed_day_and_never_today(): +async def test_first_run_rolls_up_every_closed_day_and_never_the_database_s_today(): """Before any marker exists every closed day with per-key rows is rolled up. Today is left - out: pods are still flushing it, so it is served live from the per-key table until it closes.""" + out: pods are still flushing it, so it is served live from the per-key table until it closes. + The database clock says which day that is; a pod booting with its clock a day ahead must not + roll the open day up and mark it reconciled.""" prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15")) - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14") assert result.failed_day is None @@ -114,11 +120,12 @@ async def test_first_run_rolls_up_every_closed_day_and_never_today(): @pytest.mark.asyncio async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14"), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) prisma.reconciled.clear() + prisma.today = TODAY - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-14",) assert await reconciled_through(prisma) == "2026-09-14" @@ -128,13 +135,14 @@ async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_run(): """Per-key rows carry the request start date, so a delayed flush or retry can add spend to a day far behind the marker. That day is rewritten, and the marker never moves back for it.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13"), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) prisma.reconciled.clear() + prisma.today = TODAY prisma.write_late_row("2026-09-01") prisma.write_late_row("2026-09-03") - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01", "2026-09-03") assert "2026-09-05" not in prisma.reconciled @@ -145,15 +153,16 @@ async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_ru async def test_a_late_row_seen_by_a_failed_run_is_seen_again_by_the_next_one(): """The scan time only advances when every pending day was rewritten, otherwise a late row found by the failed run would be counted as handled.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13"), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) + prisma.today = TODAY prisma.write_late_row("2026-09-01") prisma.failing_days = frozenset({"2026-09-01"}) - failed = await run_daily_global_spend_reconcile(prisma, today=TODAY) + failed = await run_daily_global_spend_reconcile(prisma) prisma.failing_days = frozenset() prisma.reconciled.clear() - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert failed.failed_day == "2026-09-01" assert failed.reconciled_through == "2026-09-13" @@ -166,7 +175,7 @@ async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-13"}' - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01", "2026-09-13") marker = await read_marker(prisma) @@ -175,11 +184,11 @@ async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): @pytest.mark.asyncio async def test_a_run_with_no_new_closed_days_keeps_the_marker(): - prisma = _FakePrisma(user_days=("2026-09-13",)) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) prisma.reconciled.clear() - result = await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == () assert result.reconciled_through == "2026-09-13" @@ -191,7 +200,7 @@ async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_goo a global table missing that day's spend.""" prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01",) assert result.failed_day == "2026-09-02" @@ -203,10 +212,10 @@ async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_goo @pytest.mark.asyncio async def test_the_next_run_resumes_from_the_failed_day(): prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) - await run_daily_global_spend_reconcile(prisma, today=TODAY) + await run_daily_global_spend_reconcile(prisma) prisma.failing_days = frozenset() - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03") assert await reconciled_through(prisma) == "2026-09-03" @@ -215,13 +224,14 @@ async def test_the_next_run_resumes_from_the_failed_day(): @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" - prisma = _FakePrisma(user_days=("2026-09-13",)) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) + prisma.today = TODAY prisma.write_late_row("2026-09-12") prisma.failing_days = frozenset({"2026-09-12"}) alert = AsyncMock() - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert, today=TODAY) + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert) assert result is not None assert result.days_reconciled == () @@ -236,7 +246,7 @@ async def test_a_clean_run_does_not_alert(): prisma = _FakePrisma(user_days=("2026-09-13",)) alert = AsyncMock() - await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert, today=TODAY) + await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert) alert.assert_not_awaited() @@ -256,7 +266,7 @@ async def test_scheduled_run_skips_when_another_pod_holds_the_lock(): prisma = _FakePrisma(user_days=("2026-09-13",)) lock = _pod_lock(acquired=False) - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) assert result is None assert prisma.reconciled == [] @@ -268,7 +278,7 @@ async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): prisma = _FakePrisma(user_days=("2026-09-13",)) lock = _pod_lock(acquired=True) - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_awaited_once() @@ -282,7 +292,7 @@ async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read() lock = _pod_lock(acquired=False) lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_not_awaited() From 0a81c6d3a8efadecc6498a13bbdd474374bb490f Mon Sep 17 00:00:00 2001 From: jesus Date: Wed, 16 Sep 2026 19:59:21 +0000 Subject: [PATCH 026/102] fix(proxy): resolve model_group_alias to its target for /v1/models metadata Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 14 +++++-- tests/test_litellm/proxy/test_proxy_utils.py | 40 ++++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 215fb143f7b..f30a3da9d68 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -192,6 +192,7 @@ from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES from litellm.types.llms.openai import ResponsesAPIResponse @@ -8177,18 +8178,23 @@ def create_model_info_response( "owned_by": provider, } - listing_info: Final = llm_router.get_model_listing_info(model_id) if llm_router is not None else None + alias_target: Final = ( + resolve_model_group_alias(llm_router.model_group_alias, model_id) if llm_router is not None else None + ) + lookup_model: Final = alias_target if alias_target is not None else model_id + + listing_info: Final = llm_router.get_model_listing_info(lookup_model) if llm_router is not None else None # One entry per distinct model behind the listed name; (None,) when the router knows # nothing about it, so the listed name is resolved on its own as before. deployment_models: Final[tuple[str | None, ...]] = ( listing_info.cost_map_keys if listing_info is not None and listing_info.cost_map_keys else (None,) ) - listed_info: Final = _safe_get_model_info(model_id, get_model_info) + listed_info: Final = _safe_get_model_info(lookup_model, get_model_info) candidate_sets: Final = tuple( _resolve_listing_model_info( deployment_model=deployment_model, - listed_model=model_id, + listed_model=lookup_model, listed_info=listed_info, get_model_info=get_model_info, ) @@ -8219,7 +8225,7 @@ def create_model_info_response( max_output_tokens = listing_info.max_output_tokens if llm_router is not None: - configured_mode: Final = llm_router.get_configured_mode(model_id) + configured_mode: Final = llm_router.get_configured_mode(lookup_model) if isinstance(configured_mode, str): base["mode"] = configured_mode diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 94ccc2762c5..a79e0798e29 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -2236,6 +2236,46 @@ def test_create_model_info_response_resolves_mode_through_deployment_model(): assert response["mode"] == "embedding" +@pytest.mark.parametrize( + "model_group_alias", + [ + {"team-embeddings": "my-embeddings"}, + {"team-embeddings": {"model": "my-embeddings", "hidden": False}}, + ], +) +def test_create_model_info_response_resolves_model_group_alias_to_target(model_group_alias): + """A `model_group_alias` row must report the metadata of the group it points at, + not the cost-map generalization or nothing that the alias name resolves to.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "my-embeddings", + "litellm_params": {"model": "openai/text-embedding-3-small"}, + } + ], + model_group_alias=model_group_alias, + ) + + alias_response = create_model_info_response( + model_id="team-embeddings", provider="openai", llm_router=router + ) + target_response = create_model_info_response( + model_id="my-embeddings", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert alias_response["id"] == "team-embeddings" + for field in ("mode", "max_input_tokens", "max_output_tokens"): + assert alias_response.get(field) == target_response.get(field) + assert alias_response["mode"] == "embedding" + + @pytest.mark.parametrize( "key_metadata, team_metadata, expected_to_run", [ From f93d80ea840e9b29f17601990a8bccec67f08f1f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 18:02:06 -0700 Subject: [PATCH 027/102] fix(mistral): forward reasoning_effort only on models that accept it --- litellm/llms/mistral/chat/transformation.py | 14 +++--- .../test_mistral_chat_transformation.py | 46 +++++++++++++++---- 2 files changed, 44 insertions(+), 16 deletions(-) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 970da0582ae..807f201a94f 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -24,7 +24,7 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.mistral import MistralThinkingBlock, MistralToolCallMessage from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, ModelResponseStream -from litellm.utils import convert_to_model_response_object +from litellm.utils import convert_to_model_response_object, supports_reasoning if TYPE_CHECKING: import tiktoken @@ -87,7 +87,9 @@ class MistralConfig(OpenAIGPTConfig): return super().get_config() def get_supported_openai_params(self, model: str) -> list[str]: - supported_params: Final = [ + is_magistral: Final = "magistral" in model.lower() + accepts_reasoning_effort: Final = is_magistral or supports_reasoning(model=model, custom_llm_provider="mistral") + return [ "stream", "temperature", "top_p", @@ -99,14 +101,10 @@ class MistralConfig(OpenAIGPTConfig): "stop", "response_format", "parallel_tool_calls", - "reasoning_effort", + *(("thinking",) if is_magistral else ()), + *(("reasoning_effort",) if accepts_reasoning_effort else ()), ] - if "magistral" in model.lower(): - supported_params.append("thinking") - - return supported_params - def _map_tool_choice(self, tool_choice: str) -> str: if tool_choice == "auto" or tool_choice == "none": return tool_choice diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py index edfaf352e1f..57a5f9ef2cd 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py @@ -51,11 +51,18 @@ class TestMistralReasoningSupport: assert "reasoning_effort" in supported_params assert "thinking" in supported_params - # Non-magistral models accept reasoning_effort (forwarded verbatim) but not thinking + # Non-magistral reasoning models accept reasoning_effort (forwarded verbatim) but not thinking + supported_params_reasoning = mistral_config.get_supported_openai_params( + "mistral/mistral-medium-latest" + ) + assert "reasoning_effort" in supported_params_reasoning + assert "thinking" not in supported_params_reasoning + + # Models Mistral rejects reasoning_effort on keep it unsupported, so drop_params still drops it supported_params_normal = mistral_config.get_supported_openai_params( "mistral/mistral-large-latest" ) - assert "reasoning_effort" in supported_params_normal + assert "reasoning_effort" not in supported_params_normal assert "thinking" not in supported_params_normal def test_map_openai_params_reasoning_effort(self): @@ -78,23 +85,46 @@ class TestMistralReasoningSupport: result_normal = mistral_config.map_openai_params( non_default_params={"reasoning_effort": "low"}, optional_params=optional_params_normal, - model="mistral/mistral-large-latest", + model="mistral/mistral-medium-latest", drop_params=False, ) assert "_add_reasoning_prompt" not in result_normal assert result_normal["reasoning_effort"] == "low" - def test_reasoning_effort_not_unsupported_for_non_magistral(self): - """Codex sends reasoning_effort to every model; Mistral must not raise UnsupportedParamsError.""" + @pytest.mark.parametrize( + ("model", "reasoning_effort"), + [("mistral-medium-latest", "high"), ("zai-glm-5-2", "xhigh")], + ) + def test_reasoning_effort_forwarded_verbatim_for_reasoning_models(self, model, reasoning_effort): + """Codex sends reasoning_effort to every model; Mistral reasoning models forward it as-is.""" import litellm optional_params = litellm.get_optional_params( - model="mistral-medium-latest", + model=model, custom_llm_provider="mistral", - reasoning_effort="medium", + reasoning_effort=reasoning_effort, ) - assert optional_params["reasoning_effort"] == "medium" + assert optional_params["reasoning_effort"] == reasoning_effort + + def test_reasoning_effort_stays_unsupported_for_non_reasoning_models(self): + """Mistral rejects reasoning_effort on codestral, so drop_params keeps dropping it there.""" + import litellm + + with pytest.raises(litellm.UnsupportedParamsError): + litellm.get_optional_params( + model="codestral-latest", + custom_llm_provider="mistral", + reasoning_effort="high", + ) + + dropped = litellm.get_optional_params( + model="codestral-latest", + custom_llm_provider="mistral", + reasoning_effort="high", + drop_params=True, + ) + assert "reasoning_effort" not in dropped def test_client_metadata_stripped_from_request(self): """client_metadata passed by Codex must not reach Mistral, whose schema rejects unknown fields.""" From b77f866dbb00b41d322d1f0dba80d49e040aa346 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 01:13:13 +0000 Subject: [PATCH 028/102] refactor(mistral): drop client_metadata without mutating optional_params Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/mistral/chat/transformation.py | 4 ++-- .../llms/mistral/test_mistral_chat_transformation.py | 6 ------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 807f201a94f..6316128e6fa 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -531,13 +531,13 @@ class MistralConfig(OpenAIGPTConfig): if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False): messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params) - optional_params.pop("client_metadata", None) + upstream_params: Final = {key: value for key, value in optional_params.items() if key != "client_metadata"} # Call parent transform_request which handles _transform_messages return super().transform_request( model=model, messages=messages, - optional_params=optional_params, + optional_params=upstream_params, litellm_params=litellm_params, headers=headers, ) diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py index 57a5f9ef2cd..38639f23050 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py @@ -51,14 +51,12 @@ class TestMistralReasoningSupport: assert "reasoning_effort" in supported_params assert "thinking" in supported_params - # Non-magistral reasoning models accept reasoning_effort (forwarded verbatim) but not thinking supported_params_reasoning = mistral_config.get_supported_openai_params( "mistral/mistral-medium-latest" ) assert "reasoning_effort" in supported_params_reasoning assert "thinking" not in supported_params_reasoning - # Models Mistral rejects reasoning_effort on keep it unsupported, so drop_params still drops it supported_params_normal = mistral_config.get_supported_openai_params( "mistral/mistral-large-latest" ) @@ -80,7 +78,6 @@ class TestMistralReasoningSupport: assert result.get("_add_reasoning_prompt") is True - # Test reasoning_effort forwarded verbatim for non-magistral model optional_params_normal = {} result_normal = mistral_config.map_openai_params( non_default_params={"reasoning_effort": "low"}, @@ -97,7 +94,6 @@ class TestMistralReasoningSupport: [("mistral-medium-latest", "high"), ("zai-glm-5-2", "xhigh")], ) def test_reasoning_effort_forwarded_verbatim_for_reasoning_models(self, model, reasoning_effort): - """Codex sends reasoning_effort to every model; Mistral reasoning models forward it as-is.""" import litellm optional_params = litellm.get_optional_params( @@ -108,7 +104,6 @@ class TestMistralReasoningSupport: assert optional_params["reasoning_effort"] == reasoning_effort def test_reasoning_effort_stays_unsupported_for_non_reasoning_models(self): - """Mistral rejects reasoning_effort on codestral, so drop_params keeps dropping it there.""" import litellm with pytest.raises(litellm.UnsupportedParamsError): @@ -127,7 +122,6 @@ class TestMistralReasoningSupport: assert "reasoning_effort" not in dropped def test_client_metadata_stripped_from_request(self): - """client_metadata passed by Codex must not reach Mistral, whose schema rejects unknown fields.""" mistral_config = MistralConfig() request = mistral_config.transform_request( From 6e1b4959d18d457f3045c97122162a37469ce975 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 03:14:13 +0000 Subject: [PATCH 029/102] feat(passthrough): deepgram streaming /v1/listen WebSocket passthrough with duration-based cost tracking Adds authenticated /deepgram/v1/listen and /deepgram/listen WebSocket routes that resolve the Deepgram credential through the pass-through router, inject Authorization: Token upstream, default the model to nova-3 when the client passes none, and relay audio and transcript frames unchanged. The shared WebSocket relay no longer assumes the first upstream frame is JSON and forwards every frame as received, keeping the Vertex AI Live setup handling on Vertex routes only. A Deepgram logging handler bills the call on Metadata.duration, falling back to the furthest Results start + duration, at the deepgram/ per-second rate from the model cost map Resolves LIT-7937 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 3 + litellm/llms/deepgram/common_utils.py | 22 ++ litellm/proxy/_lazy_features.py | 1 + litellm/proxy/_types.py | 3 + .../llm_passthrough_endpoints.py | 68 ++++- ...gram_listen_passthrough_logging_handler.py | 132 +++++++++ .../pass_through_endpoints.py | 113 ++++--- .../pass_through_endpoints/success_handler.py | 18 ++ ...gram_listen_passthrough_logging_handler.py | 254 ++++++++++++++++ .../test_deepgram_ws_passthrough_routes.py | 280 ++++++++++++++++++ .../test_pass_through_endpoints.py | 163 +++++++++- 11 files changed, 974 insertions(+), 83 deletions(-) create mode 100644 litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py diff --git a/litellm/constants.py b/litellm/constants.py index 8409a161800..c3a7a16ea39 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -317,6 +317,9 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float( # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 +DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1" +DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3" + BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index a741b092a36..db00d048f01 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -1,5 +1,27 @@ +from types import MappingProxyType +from typing import Final + +import httpx + +from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL from litellm.llms.base_llm.chat.transformation import BaseLLMException +_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"}) + class DeepgramException(BaseLLMException): pass + + +def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str: + """ + The upstream ``/listen`` socket for a streaming transcription, keeping the client's query string as sent + and adding the default model only when the client named none + """ + listen_url: Final = httpx.URL(f"{(api_base or DEEPGRAM_DEFAULT_API_BASE).rstrip('/')}/listen") + websocket_url: Final = listen_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(listen_url.scheme, listen_url.scheme)) + params: Final = httpx.QueryParams(query_string) + query: Final = ( + query_string if params.get("model") else str(params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL)) + ) + return f"{websocket_url}?{query}" diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index faf95397fa5..ce48c4801d6 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -200,6 +200,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( "/cohere/", "/comprehendmedical", "/cursor/", + "/deepgram/", "/eu.assemblyai/", "/gemini/", "/gigachat/", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 680c63393e8..aa2068cb94e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -70,6 +70,7 @@ from litellm.types.utils import ( StandardLoggingVectorStoreRequest, StandardPassThroughResponseObject, TextCompletionResponse, + TranscriptionResponse, ) from litellm.types.videos.main import VideoObject @@ -487,6 +488,7 @@ class LiteLLMRoutes(enum.Enum): "/gigachat", "/watsonx", "/nvidia_nim", + "/deepgram", ] ######################################################### @@ -4694,6 +4696,7 @@ PassThroughEndpointLoggingResultValues = ( | VideoObject | StandardPassThroughResponseObject | ResponsesAPIResponse + | TranscriptionResponse ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index b9b8cb3a22b..a5a95c34910 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -36,6 +36,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.deepgram.common_utils import deepgram_listen_websocket_target from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse @@ -2573,7 +2574,7 @@ async def _openai_websocket_refusal( return None -class _OpenAIWebsocketRelay(Protocol): +class _WebsocketRelay(Protocol): async def __call__( self, *, @@ -2593,7 +2594,7 @@ def _proxy_general_settings() -> Mapping[str, object]: return general_settings -def _openai_websocket_relay() -> _OpenAIWebsocketRelay: +def _websocket_relay() -> _WebsocketRelay: return websocket_passthrough_request @@ -2611,6 +2612,19 @@ def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists: return resolve +def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: + """ + The first subprotocol the client offered, echoed back so browsers that carry the LiteLLM key in + ``Sec-WebSocket-Protocol`` complete the handshake + """ + requested_subprotocols: Final = tuple( + protocol.strip() + for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") + if protocol.strip() + ) + return requested_subprotocols[0] if requested_subprotocols else None + + @router.websocket("/openai_passthrough/{endpoint:path}") @router.websocket("/openai/{endpoint:path}") async def openai_websocket_proxy_route( @@ -2618,16 +2632,11 @@ async def openai_websocket_proxy_route( endpoint: str, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], - relay: Annotated[_OpenAIWebsocketRelay, Depends(_openai_websocket_relay)], + relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" - requested_subprotocols: Final = tuple( - protocol.strip() - for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") - if protocol.strip() - ) - negotiated_subprotocol: Final = requested_subprotocols[0] if requested_subprotocols else None + negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket) refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists) if refusal is not None: @@ -2686,6 +2695,47 @@ async def openai_websocket_proxy_route( ) +_DEEPGRAM_WS_MISSING_KEY_REASON: Final = ( + "Required 'DEEPGRAM_API_KEY' in environment to make pass-through calls to Deepgram." +) + + +@router.websocket("/deepgram/v1/listen") +@router.websocket("/deepgram/listen") +async def deepgram_listen_websocket_route( + websocket: WebSocket, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], + relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], +) -> None: + """ + Streaming speech to text through Deepgram's ``/v1/listen`` socket. Audio frames and transcript frames are + relayed unchanged; the call is billed on the audio duration Deepgram reports when the socket closes + """ + deepgram_api_key: Final = passthrough_endpoint_router.get_credentials( + custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value, + region_name=None, + ) + if deepgram_api_key is None: + await websocket.close(code=1011, reason=_DEEPGRAM_WS_MISSING_KEY_REASON) + return + + await websocket.accept(subprotocol=_negotiated_websocket_subprotocol(websocket)) + await relay( + websocket=websocket, + target=deepgram_listen_websocket_target( + api_base=get_secret_str("DEEPGRAM_API_BASE"), + query_string=websocket.url.query, + ), + custom_headers={ # mutable-ok: websocket_passthrough_request requires a plain dict of upstream headers + "Authorization": f"Token {deepgram_api_key}" + }, + user_api_key_dict=user_api_key_dict, + forward_headers=False, + endpoint=websocket.url.path, + accept_websocket=False, + ) + + class BaseOpenAIPassThroughHandler: @staticmethod async def _base_openai_pass_through_handler( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py new file mode 100644 index 00000000000..6a574a2e1b9 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py @@ -0,0 +1,132 @@ +""" +Cost tracking for Deepgram's streaming ``/v1/listen`` WebSocket. Deepgram bills the audio it processed, which it +reports as ``duration`` on the closing ``Metadata`` frame; a stream that ends without one is billed on the furthest +``start + duration`` across its ``Results`` frames +""" + +import math +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final +from urllib.parse import parse_qs, urlparse + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.constants import DEEPGRAM_LISTEN_DEFAULT_MODEL +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.types.utils import TranscriptionResponse + +DEEPGRAM_LISTEN_ROUTE_SUFFIX: Final = "/listen" + + +def _seconds(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) if math.isfinite(value) and value >= 0 else None + + +def _results_frame_end(frame: Mapping[str, object]) -> float | None: + start: Final = _seconds(frame.get("start")) + duration: Final = _seconds(frame.get("duration")) + return None if start is None or duration is None else start + duration + + +def _final_transcript(frame: Mapping[str, object]) -> str | None: + if frame.get("is_final") is not True: + return None + channel: Final = frame.get("channel") + alternatives: Final = channel.get("alternatives") if isinstance(channel, Mapping) else None + first: Final = alternatives[0] if isinstance(alternatives, list) and alternatives else None + transcript: Final = first.get("transcript") if isinstance(first, Mapping) else None + return transcript if isinstance(transcript, str) and transcript else None + + +def deepgram_listen_audio_seconds(websocket_messages: Sequence[Mapping[str, object]]) -> float: + metadata_durations: Final = tuple( + duration + for frame in websocket_messages + if frame.get("type") == "Metadata" + if (duration := _seconds(frame.get("duration"))) is not None + ) + if metadata_durations: + return metadata_durations[-1] + return max( + ( + end + for frame in websocket_messages + if frame.get("type") == "Results" + if (end := _results_frame_end(frame)) is not None + ), + default=0.0, + ) + + +def deepgram_listen_transcript(websocket_messages: Sequence[Mapping[str, object]]) -> str: + return " ".join( + transcript + for frame in websocket_messages + if frame.get("type") == "Results" + if (transcript := _final_transcript(frame)) is not None + ) + + +def deepgram_listen_model(upstream_url: str) -> str: + models: Final = parse_qs(urlparse(upstream_url).query).get("model") + return models[0] if models else DEEPGRAM_LISTEN_DEFAULT_MODEL + + +def _audio_cost(response: TranscriptionResponse, model: str) -> float | None: + try: + return litellm.completion_cost( + completion_response=response, + model=model, + custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value, + call_type="transcription", + ) + except Exception as e: # noqa: BLE001 # an unpriced model must not lose the spend row, only its cost + verbose_proxy_logger.warning("Deepgram listen passthrough: no pricing for model '%s': %s", model, e) + return None + + +class DeepgramListenPassthroughLoggingHandler: + @staticmethod + def is_deepgram_listen_route(url_route: str) -> bool: + path: Final = urlparse(url_route).path + return path.startswith("/deepgram/") and path.endswith(DEEPGRAM_LISTEN_ROUTE_SUFFIX) + + def deepgram_listen_passthrough_handler( + self, + websocket_messages: Sequence[Mapping[str, object]], + logging_obj: LiteLLMLoggingObj, + upstream_url: str, + kwargs: Mapping[str, object] = MappingProxyType({}), + ) -> PassThroughEndpointLoggingTypedDict: + model: Final = deepgram_listen_model(upstream_url) + audio_seconds: Final = deepgram_listen_audio_seconds(websocket_messages) + response: Final = TranscriptionResponse(text=deepgram_listen_transcript(websocket_messages)) + response._hidden_params["audio_transcription_duration"] = audio_seconds # pyright: ignore[reportPrivateUsage] # the cost calculator reads the billed duration off the response's hidden params + response_cost: Final = _audio_cost(response, model) + response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads a precomputed cost off the response's hidden params + + provider: Final = litellm.LlmProviders.DEEPGRAM.value + logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object + logging_obj.model_call_details["model"] = model # rebind-ok: same shared logging object + logging_obj.model_call_details["custom_llm_provider"] = provider # rebind-ok: same shared logging object + logging_obj.model_call_details["response_cost"] = response_cost # rebind-ok: same shared logging object + verbose_proxy_logger.debug( + "Deepgram listen passthrough cost tracking: model %s, audio seconds %s, cost %s", + model, + audio_seconds, + response_cost, + ) + logging_result: Final[PassThroughEndpointLoggingTypedDict] = { + "result": response, + "kwargs": { + **kwargs, + "model": model, + "custom_llm_provider": provider, + "response_cost": response_cost, + }, + } + return logging_result diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 685c19062bb..cf6985852f2 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -8,7 +8,7 @@ from base64 import b64encode from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime -from itertools import groupby +from itertools import count, groupby from typing import TYPE_CHECKING, Any, Final, TypedDict, cast from urllib.parse import urlencode, urlparse @@ -2120,6 +2120,17 @@ def _resolved_vertex_live_setup( return {**setup_data, "model": setup_model_rewriter(setup_model)} +def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: + """ + The frame as a JSON object when it is one, for cost tracking; audio and non-object frames yield None + """ + try: + decoded: Final = json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + return None + return decoded if isinstance(decoded, dict) else None + + def _truncated_close_reason(reason: str) -> str: """ Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character @@ -2401,70 +2412,46 @@ async def websocket_passthrough_request( ) await upstream_ws.close() + def _extract_vertex_live_model_from_setup_response(setup_response: Mapping[str, object]) -> None: + extracted_model: Final = _extract_model_from_vertex_ai_setup(setup_response) + if not extracted_model: + verbose_proxy_logger.warning( + "WebSocket passthrough (%s): Failed to extract model from server setup response: %s", + endpoint, + setup_response, + ) + return + kwargs["model"] = extracted_model + kwargs["custom_llm_provider"] = "vertex_ai_language_models" + logging_obj.model = extracted_model + logging_obj.model_call_details["model"] = extracted_model + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models" + + is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint) + json_frame_ordinal: Final = count() + + async def relay_upstream_frame(upstream_message: str | bytes) -> None: + """ + Send the frame to the client exactly as received, then keep it for cost tracking when it is a JSON + object; the Vertex AI Live setup acknowledgement only names the model, so it is read instead of kept + """ + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + else: + await websocket.send_text(upstream_message) + message_data: Final = _json_object_frame(upstream_message) + if message_data is None: + return + if is_vertex_live and next(json_frame_ordinal) == 0: + _extract_vertex_live_model_from_setup_response(message_data) + return + websocket_messages.append(message_data) + async def forward_upstream_to_client() -> Close | None: - """Forward messages from upstream to client WebSocket, returning the upstream's close frame""" + """Relay upstream frames to the client until the upstream closes, returning its close frame""" try: - # Wait for the first response from upstream - raw_response = await upstream_ws.recv(decode=False) - # Ensure raw_response is bytes before decoding - if isinstance(raw_response, str): - raw_response = raw_response.encode("utf-8") - setup_response: Final[Mapping[str, object]] = json.loads(raw_response.decode("utf-8")) - verbose_proxy_logger.debug("Setup response: %s", setup_response) - - # Extract model and provider from setup response for Vertex AI Live - if endpoint and "/vertex_ai/live" in endpoint: - verbose_proxy_logger.debug( - "WebSocket passthrough (%s): Processing server setup response for model extraction", - endpoint, - ) - extracted_model: Final = _extract_model_from_vertex_ai_setup(setup_response) - if extracted_model: - kwargs["model"] = extracted_model - kwargs["custom_llm_provider"] = "vertex_ai_language_models" - # Update logging object with correct model - logging_obj.model = extracted_model - logging_obj.model_call_details["model"] = extracted_model - logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models" - verbose_proxy_logger.debug( - "WebSocket passthrough (%s): Successfully extracted model '%s' and set provider to 'vertex_ai' from server setup response", - endpoint, - extracted_model, - ) - else: - verbose_proxy_logger.warning( - "WebSocket passthrough (%s): Failed to extract model from server setup response: %s", - endpoint, - setup_response, - ) - else: - verbose_proxy_logger.debug( - "WebSocket passthrough (%s): Not a Vertex AI Live endpoint, skipping model extraction", - endpoint, - ) - - # Send the setup response to the client - await websocket.send_text(json.dumps(setup_response)) - - # Now continuously forward messages from upstream to client - async for upstream_message in upstream_ws: - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - # Parse and collect for cost tracking - try: - message_data: dict[str, object] = json.loads(upstream_message.decode()) - websocket_messages.append(message_data) - except (json.JSONDecodeError, UnicodeDecodeError): - pass - else: - await websocket.send_text(upstream_message) - # Parse and collect for cost tracking - try: - message_data = json.loads(upstream_message) - websocket_messages.append(message_data) - except json.JSONDecodeError: - pass - + while True: + await relay_upstream_frame(await upstream_ws.recv()) except (ConnectionClosedOK, ConnectionClosedError) as e: verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e) return e.rcvd diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 76a471302f4..82bb47a60ab 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -24,6 +24,9 @@ from .llm_provider_handlers.cohere_passthrough_logging_handler import ( from .llm_provider_handlers.cursor_passthrough_logging_handler import ( CursorPassthroughLoggingHandler, ) +from .llm_provider_handlers.deepgram_listen_passthrough_logging_handler import ( + DeepgramListenPassthroughLoggingHandler, +) from .llm_provider_handlers.gemini_passthrough_logging_handler import ( GeminiPassthroughLoggingHandler, ) @@ -278,6 +281,21 @@ class PassThroughEndpointLogging: standard_logging_response_object = vertex_ai_live_handler_result["result"] kwargs = vertex_ai_live_handler_result["kwargs"] + elif DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route): + deepgram_handler_result: Final = ( + DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=tuple( + message + for message in (response_body if isinstance(response_body, list) else ()) + if isinstance(message, dict) + ), + logging_obj=logging_obj, + upstream_url=str(httpx_response.request.url), + kwargs=kwargs, + ) + ) + standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain + kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract return_dict["standard_logging_response_object"] = standard_logging_response_object return_dict["kwargs"] = kwargs diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py new file mode 100644 index 00000000000..f00411ac10c --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py @@ -0,0 +1,254 @@ +"""Deepgram ``/v1/listen`` WebSocket passthrough: duration extraction and duration based cost tracking.""" + +import math +from collections.abc import Mapping, Sequence +from datetime import datetime +from types import SimpleNamespace +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.deepgram_listen_passthrough_logging_handler import ( + DeepgramListenPassthroughLoggingHandler, + deepgram_listen_audio_seconds, + deepgram_listen_model, + deepgram_listen_transcript, +) +from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging +from litellm.types.passthrough_endpoints.pass_through_endpoints import PassthroughStandardLoggingPayload +from litellm.types.utils import StandardLoggingPayload, TranscriptionResponse + +NOVA_3_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" + + +def _results(start: object, duration: object, transcript: str = "", is_final: object = True) -> dict[str, object]: + return { + "type": "Results", + "start": start, + "duration": duration, + "is_final": is_final, + "channel": {"alternatives": [{"transcript": transcript, "confidence": 0.9}]}, + } + + +def _metadata(duration: object) -> dict[str, object]: + return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": 1} + + +@pytest.mark.parametrize( + ("frames", "expected_seconds"), + [ + pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _metadata(6.25)), 6.25, id="metadata wins"), + pytest.param((_metadata(4.0), _results(0.0, 9.0), _metadata(5.5)), 5.5, id="last metadata wins"), + pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _results(1.0, 1.0)), 5.5, id="furthest results end"), + pytest.param((_results(0.0, 0.0), _metadata(0.0)), 0.0, id="zero metadata is a real zero"), + pytest.param((_results(0.0, 1.5), _metadata("6.25")), 1.5, id="string metadata is ignored"), + pytest.param((_results(0.0, 1.5), _metadata(True)), 1.5, id="boolean metadata is ignored"), + pytest.param((_results(0.0, 1.5), _metadata(-3.0)), 1.5, id="negative metadata is ignored"), + pytest.param((_results(0.0, 1.5), _metadata(math.nan), _metadata(math.inf)), 1.5, id="nan/inf ignored"), + pytest.param((_results("0", 2.0), _results(0.0, None), _results(0.0, 0.75)), 0.75, id="malformed results"), + pytest.param(({"type": "SpeechStarted", "timestamp": 3.0}, {"type": "UtteranceEnd"}), 0.0, id="no usage"), + pytest.param((), 0.0, id="no frames"), + ], +) +def test_deepgram_listen_audio_seconds(frames: Sequence[Mapping[str, object]], expected_seconds: float): + assert deepgram_listen_audio_seconds(frames) == expected_seconds + + +def test_deepgram_listen_transcript_joins_final_results_only(): + frames = ( + _results(0.0, 1.0, "hello wor", is_final=False), + _results(0.0, 1.5, "hello world"), + _results(1.5, 0.5, "", is_final=True), + _results(2.0, 1.0, "how are you", is_final="yes"), + {"type": "Results", "start": 3.0, "duration": 1.0, "is_final": True, "channel": {"alternatives": []}}, + _results(4.0, 1.0, "goodbye"), + _metadata(5.0), + ) + assert deepgram_listen_transcript(frames) == "hello world goodbye" + + +@pytest.mark.parametrize( + ("upstream_url", "expected_model"), + [ + (NOVA_3_URL, "nova-3"), + ("wss://api.deepgram.com/v1/listen?encoding=linear16&model=nova-2-medical", "nova-2-medical"), + ("wss://api.deepgram.com/v1/listen?model=nova-3&model=nova-2", "nova-3"), + ("wss://api.deepgram.com/v1/listen?encoding=linear16", litellm.constants.DEEPGRAM_LISTEN_DEFAULT_MODEL), + ], +) +def test_deepgram_listen_model_comes_from_the_upstream_query(upstream_url: str, expected_model: str): + assert deepgram_listen_model(upstream_url) == expected_model + + +@pytest.mark.parametrize( + ("url_route", "expected"), + [ + ("/deepgram/v1/listen", True), + ("/deepgram/listen", True), + ("/deepgram/v1/listen?model=nova-3", True), + ("/deepgram/v1/speak", False), + ("/deepgram/v1/listen/extra", False), + ("/openai/v1/realtime", False), + ("/vertex_ai/live", False), + ("", False), + ], +) +def test_is_deepgram_listen_route(url_route: str, expected: bool): + assert DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route) is expected + + +def _logging_obj(call_id: str = "call-dg") -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model="unknown", + messages=[{"role": "user", "content": "WebSocket connection"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id="websocket_passthrough", + ) + + +def _registry_cost(model: str, seconds: float) -> float: + """Derives the expected charge from the live cost map rather than pinning a vendor price.""" + per_second: Final = litellm.model_cost[f"deepgram/{model}"]["input_cost_per_second"] + assert per_second > 0 + return per_second * seconds + + +def test_handler_bills_metadata_duration_at_the_registry_rate_and_names_the_model(): + frames = (_results(0.0, 5.0, "first sentence"), _results(5.0, 7.5, "second sentence"), _metadata(12.5)) + logging_obj = _logging_obj() + + handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=frames, + logging_obj=logging_obj, + upstream_url=NOVA_3_URL, + kwargs={"litellm_params": {"metadata": {}}}, + ) + + result = handler_result["result"] + assert isinstance(result, TranscriptionResponse) + assert result.text == "first sentence second sentence" + assert result._hidden_params["audio_transcription_duration"] == 12.5 + assert result._hidden_params["response_cost"] == pytest.approx(_registry_cost("nova-3", 12.5)) + assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-3", 12.5)) + assert handler_result["kwargs"]["model"] == "nova-3" + assert handler_result["kwargs"]["custom_llm_provider"] == "deepgram" + assert handler_result["kwargs"]["litellm_params"] == {"metadata": {}} + assert logging_obj.model == "nova-3" + assert logging_obj.model_call_details["model"] == "nova-3" + assert logging_obj.model_call_details["custom_llm_provider"] == "deepgram" + assert logging_obj.model_call_details["response_cost"] == pytest.approx(_registry_cost("nova-3", 12.5)) + + +def test_handler_falls_back_to_results_frames_when_the_stream_ends_without_metadata(): + frames = (_results(0.0, 30.0, "a"), _results(30.0, 30.0, "b"), _results(60.0, 12.5, "c")) + + handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=frames, logging_obj=_logging_obj(), upstream_url=NOVA_3_URL + ) + + assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 72.5 + assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-3", 72.5)) + + +def test_handler_charges_more_for_more_audio_on_the_same_model(): + short = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=(_metadata(10.0),), logging_obj=_logging_obj(), upstream_url=NOVA_3_URL + ) + long = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=(_metadata(30.0),), logging_obj=_logging_obj(), upstream_url=NOVA_3_URL + ) + + assert long["kwargs"]["response_cost"] == pytest.approx(3 * short["kwargs"]["response_cost"]) + assert short["kwargs"]["response_cost"] > 0 + + +def test_handler_keeps_the_spend_row_but_no_cost_for_an_unpriced_model(): + handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=(_metadata(12.5),), + logging_obj=_logging_obj(), + upstream_url="wss://api.deepgram.com/v1/listen?model=nova-99-not-in-registry", + ) + + assert handler_result["kwargs"]["model"] == "nova-99-not-in-registry" + assert handler_result["kwargs"]["response_cost"] is None + assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 12.5 + + +class _CapturingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.payloads: list[StandardLoggingPayload] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payloads.append(kwargs["standard_logging_object"]) + + +@pytest.mark.asyncio +async def test_success_handler_dispatches_deepgram_listen_and_logs_duration_based_spend(monkeypatch): + """Drives the shared passthrough success handler the way the WebSocket relay does at socket close and reads + what a spend logger receives: Deepgram model and provider, the audio duration billed at the registry rate.""" + capturing_logger = _CapturingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [capturing_logger]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + logging_obj = _logging_obj("call-dg-e2e") + frames = [_results(0.0, 5.0, "hello world", is_final=False), _results(0.0, 5.0, "hello world"), _metadata(20.0)] + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", team_id="team-stt", user_id="user-1") + start_time = datetime.now() + passthrough_logging_payload = PassthroughStandardLoggingPayload( + url=NOVA_3_URL, request_body={}, request_method="WEBSOCKET", cost_per_request=None + ) + logging_obj.update_environment_variables( + model="unknown", + user="unknown", + optional_params={}, + litellm_params={ + "metadata": { + "user_api_key": user_api_key_dict.api_key, + "user_api_key_team_id": user_api_key_dict.team_id, + "user_api_key_user_id": user_api_key_dict.user_id, + } + }, + call_type="pass_through_endpoint", + ) + + await PassThroughEndpointLogging().pass_through_async_success_handler( + httpx_response=SimpleNamespace( + status_code=200, + text="WebSocket connection successful", + headers={}, + request=SimpleNamespace(method="WEBSOCKET", url=NOVA_3_URL), + ), + response_body=frames, + logging_obj=logging_obj, + url_route="/deepgram/v1/listen", + result="websocket_connection_successful", + start_time=start_time, + end_time=datetime.now(), + cache_hit=False, + request_body={}, + passthrough_logging_payload=passthrough_logging_payload, + litellm_params={ + "metadata": { + "user_api_key": user_api_key_dict.api_key, + "user_api_key_team_id": user_api_key_dict.team_id, + "user_api_key_user_id": user_api_key_dict.user_id, + } + }, + ) + + assert len(capturing_logger.payloads) == 1 + payload = capturing_logger.payloads[0] + assert payload["model"] == "nova-3" + assert payload["custom_llm_provider"] == "deepgram" + assert payload["response_cost"] == pytest.approx(_registry_cost("nova-3", 20.0)) + assert payload["metadata"]["user_api_key_team_id"] == "team-stt" + assert payload["id"] == "call-dg-e2e" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py new file mode 100644 index 00000000000..64804e621ba --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -0,0 +1,280 @@ +"""Deepgram ``/v1/listen`` passthrough WebSocket route: registration, auth, credential injection, target URL.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType, SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from starlette.routing import WebSocketRoute +from starlette.websockets import WebSocketDisconnect + +from litellm.proxy._lazy_features import LAZY_FEATURES +from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _websocket_relay, + deepgram_listen_websocket_route, + router, +) + +GET_CREDENTIALS: Final = ( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials" +) +USER_API_KEY_AUTH: Final = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" +LISTEN_PATHS: Final = ("/deepgram/v1/listen", "/deepgram/listen") + + +class _FakeWebSocket: + def __init__(self, path: str, query: str) -> None: + self.url = SimpleNamespace(path=path, query=query) + self.headers = {"authorization": "Bearer sk-litellm-virtual", "x-api-key": "sk-caller-secret"} + self.accepts: list[str | None] = [] + self.closed: tuple[int, str] | None = None + + async def accept(self, subprotocol: str | None = None) -> None: + self.accepts.append(subprotocol) + + async def close(self, code: int = 1000, reason: str = "") -> None: + self.closed = (code, reason) + + +@dataclass(frozen=True, slots=True) +class _RelayCall: + target: str + custom_headers: Mapping[str, str] + user_api_key_dict: UserAPIKeyAuth + forward_headers: bool + endpoint: str + accept_websocket: bool + + +class _FakeRelay: + def __init__(self) -> None: + self.calls: list[_RelayCall] = [] + + async def __call__( + self, + *, + websocket: object, + target: str, + custom_headers: dict[str, str], + user_api_key_dict: UserAPIKeyAuth, + forward_headers: bool, + endpoint: str, + accept_websocket: bool, + ) -> None: + self.calls.append( + _RelayCall( + target=target, + custom_headers=MappingProxyType(dict(custom_headers)), + user_api_key_dict=user_api_key_dict, + forward_headers=forward_headers, + endpoint=endpoint, + accept_websocket=accept_websocket, + ) + ) + + +async def _serve(websocket: _FakeWebSocket, user_api_key_dict: UserAPIKeyAuth | None = None) -> _FakeRelay: + relay = _FakeRelay() + await deepgram_listen_websocket_route( + websocket=websocket, + user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(), + relay=relay, + ) + return relay + + +def test_deepgram_listen_websocket_routes_registered(): + ws_paths = {route.path for route in router.routes if isinstance(route, WebSocketRoute)} + assert set(LISTEN_PATHS) <= ws_paths + + +@pytest.mark.parametrize("path", LISTEN_PATHS) +def test_deepgram_listen_is_a_lazily_loaded_mapped_pass_through_route(path): + """The route must be reachable before the passthrough module is imported and must be authed and + billed as a mapped pass-through route like the other provider prefixes.""" + feature = next(feature for feature in LAZY_FEATURES if feature.name == "llm_passthrough") + assert feature.matches(path) + assert any(path.startswith(prefix) for prefix in LiteLLMRoutes.mapped_pass_through_routes.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", LISTEN_PATHS) +async def test_deepgram_listen_forwards_query_and_injects_only_provider_auth(path, monkeypatch): + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + websocket = _FakeWebSocket(path, "encoding=linear16&sample_rate=16000&keywords=hi%3A2&keywords=there") + caller = UserAPIKeyAuth(api_key="sk-litellm-virtual", team_id="team-stt") + + with patch(GET_CREDENTIALS, return_value="dg-provider-key") as get_credentials: + relay = await _serve(websocket, caller) + + assert get_credentials.call_args.kwargs == {"custom_llm_provider": "deepgram", "region_name": None} + assert relay.calls == [ + _RelayCall( + target=( + "wss://api.deepgram.com/v1/listen" + "?encoding=linear16&sample_rate=16000&keywords=hi%3A2&keywords=there&model=nova-3" + ), + custom_headers=MappingProxyType({"Authorization": "Token dg-provider-key"}), + user_api_key_dict=caller, + forward_headers=False, + endpoint=path, + accept_websocket=False, + ) + ] + assert websocket.accepts == [None] + assert websocket.closed is None + + +@pytest.mark.asyncio +async def test_deepgram_listen_keeps_caller_chosen_model(monkeypatch): + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-2&language=en") + + with patch(GET_CREDENTIALS, return_value="dg-provider-key"): + relay = await _serve(websocket) + + assert [call.target for call in relay.calls] == ["wss://api.deepgram.com/v1/listen?model=nova-2&language=en"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("query", "expected_target"), + [ + ("", "wss://api.deepgram.com/v1/listen?model=nova-3"), + ("model=", "wss://api.deepgram.com/v1/listen?model=nova-3"), + ("model=&language=en", "wss://api.deepgram.com/v1/listen?language=en&model=nova-3"), + ], +) +async def test_deepgram_listen_defaults_to_nova_3_when_no_model_is_named(query, expected_target, monkeypatch): + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + websocket = _FakeWebSocket("/deepgram/listen", query) + + with patch(GET_CREDENTIALS, return_value="dg-provider-key"): + relay = await _serve(websocket) + + assert [call.target for call in relay.calls] == [expected_target] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("api_base", "expected_target"), + [ + ("https://api.eu.deepgram.com/v1/", "wss://api.eu.deepgram.com/v1/listen?model=nova-3"), + ("http://localhost:8080/v1", "ws://localhost:8080/v1/listen?model=nova-3"), + ("wss://deepgram.internal.example/v1", "wss://deepgram.internal.example/v1/listen?model=nova-3"), + ], +) +async def test_deepgram_listen_honours_server_configured_api_base(api_base, expected_target, monkeypatch): + monkeypatch.setenv("DEEPGRAM_API_BASE", api_base) + websocket = _FakeWebSocket("/deepgram/v1/listen", "") + + with patch(GET_CREDENTIALS, return_value="dg-provider-key"): + relay = await _serve(websocket) + + assert [call.target for call in relay.calls] == [expected_target] + + +@pytest.mark.asyncio +async def test_deepgram_listen_ignores_caller_supplied_api_base(monkeypatch): + """V1: the server-configured Deepgram key must only ever go to the server-configured host.""" + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + websocket = _FakeWebSocket("/deepgram/v1/listen", "api_base=wss%3A%2F%2Fattacker.example%2Fv1&model=nova-3") + + with patch(GET_CREDENTIALS, return_value="dg-provider-key"): + relay = await _serve(websocket) + + assert [call.target for call in relay.calls] == [ + "wss://api.deepgram.com/v1/listen?api_base=wss%3A%2F%2Fattacker.example%2Fv1&model=nova-3" + ] + + +@pytest.mark.asyncio +async def test_deepgram_listen_closes_cleanly_when_provider_credentials_missing(): + websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-3") + + with patch(GET_CREDENTIALS, return_value=None): + relay = await _serve(websocket) + + assert websocket.closed is not None + assert websocket.closed[0] == 1011 + assert "DEEPGRAM_API_KEY" in websocket.closed[1] + assert websocket.accepts == [] + assert relay.calls == [] + + +def _app_with_relay(relay: _FakeRelay) -> FastAPI: + app = FastAPI() + app.include_router(router) + app.dependency_overrides[_websocket_relay] = lambda: relay + return app + + +def test_deepgram_listen_rejects_connections_without_a_litellm_key(): + relay = _FakeRelay() + client = TestClient(_app_with_relay(relay)) + + with patch(GET_CREDENTIALS, return_value="dg-provider-key") as get_credentials: + with pytest.raises(WebSocketDisconnect) as disconnect: + with client.websocket_connect("/deepgram/v1/listen?model=nova-3"): + pass + + assert disconnect.value.code == 1008 + assert relay.calls == [] + get_credentials.assert_not_called() + + +def test_deepgram_listen_authenticates_the_litellm_key_and_relays_to_deepgram(monkeypatch): + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + relay = _FakeRelay() + client = TestClient(_app_with_relay(relay)) + caller = UserAPIKeyAuth(api_key="hashed-sk-litellm", team_id="team-stt") + + with ( + patch(GET_CREDENTIALS, return_value="dg-provider-key"), + patch(USER_API_KEY_AUTH, new=AsyncMock(return_value=caller)) as auth, + ): + with client.websocket_connect( + "/deepgram/v1/listen?model=nova-3&punctuate=true", + headers={"Authorization": "Bearer sk-litellm-virtual"}, + ): + pass + + assert auth.await_args.kwargs["api_key"] == "Bearer sk-litellm-virtual" + assert relay.calls == [ + _RelayCall( + target="wss://api.deepgram.com/v1/listen?model=nova-3&punctuate=true", + custom_headers=MappingProxyType({"Authorization": "Token dg-provider-key"}), + user_api_key_dict=caller, + forward_headers=False, + endpoint="/deepgram/v1/listen", + accept_websocket=False, + ) + ] + + +def test_deepgram_listen_echoes_the_browser_subprotocol_that_carries_the_litellm_key(monkeypatch): + """Browsers cannot set headers, so they send the key as a subprotocol and abort the handshake unless the + server echoes that subprotocol back; the key itself must still stay off the upstream connection.""" + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + relay = _FakeRelay() + client = TestClient(_app_with_relay(relay)) + + with ( + patch(GET_CREDENTIALS, return_value="dg-provider-key"), + patch(USER_API_KEY_AUTH, new=AsyncMock(return_value=UserAPIKeyAuth(api_key="hashed"))), + ): + with client.websocket_connect( + "/deepgram/v1/listen?model=nova-3", + subprotocols=["openai-insecure-api-key.sk-litellm-virtual"], + ) as connection: + assert connection.accepted_subprotocol == "openai-insecure-api-key.sk-litellm-virtual" + + assert [call.custom_headers for call in relay.calls] == [ + MappingProxyType({"Authorization": "Token dg-provider-key"}) + ] + assert [call.forward_headers for call in relay.calls] == [False] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d854ee39ff4..de13e52498f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4844,18 +4844,21 @@ async def test_unusable_upstream_cost_records_zero_not_the_flat_estimate(): class FakeUpstreamWebSocket: - def __init__(self, first_frame: bytes): - self._first_frame = first_frame + """Serves the given frames in order, then closes normally, the way a real websockets connection does""" + + def __init__(self, *frames: str | bytes): + self._frames = iter(frames) self.close = AsyncMock() + self.send = AsyncMock() - async def recv(self, decode: bool = True): - return self._first_frame + async def recv(self, decode: bool | None = None): + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close - def __aiter__(self): - return self - - async def __anext__(self): - raise StopAsyncIteration + frame = next(self._frames, None) + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame class FakeUpstreamConnect: @@ -4876,7 +4879,7 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): first_frame = json.dumps( {"type": "session.created", "session": {"instructions": "Hablas español, ¿sí?"}}, ensure_ascii=False, - ).encode("utf-8") + ) upstream_ws = FakeUpstreamWebSocket(first_frame) websocket = MagicMock() @@ -4930,7 +4933,7 @@ async def test_websocket_passthrough_propagates_active_trace_context( from starlette.websockets import WebSocketState captured: dict[str, dict[str, str]] = {} - upstream_ws = FakeUpstreamWebSocket(b"{}") + upstream_ws = FakeUpstreamWebSocket("{}") def fake_connect(target, additional_headers): captured["headers"] = additional_headers @@ -5359,6 +5362,144 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason) +DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" +DEEPGRAM_INTERIM_FRAME = json.dumps( + { + "type": "Results", + "start": 0.0, + "duration": 1.02, + "is_final": False, + "channel": {"alternatives": [{"transcript": "hello wor", "confidence": 0.71}]}, + } +) +DEEPGRAM_FINAL_FRAME = json.dumps( + { + "type": "Results", + "start": 0.0, + "duration": 2.5, + "is_final": True, + "speech_final": True, + "channel": {"alternatives": [{"transcript": "hello world, ¿qué tal?", "confidence": 0.98}]}, + }, + ensure_ascii=False, +) +DEEPGRAM_METADATA_FRAME = json.dumps({"type": "Metadata", "request_id": "req-1", "duration": 2.5, "channels": 1}) + + +async def _relay_deepgram_listen(upstream_ws, client_receive): + """Runs the generic relay the way the Deepgram route does and returns (client websocket, success handler mock)""" + websocket = _client_websocket(client_receive) + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target=DEEPGRAM_LISTEN_TARGET, + custom_headers={"Authorization": "Token dg-provider-key"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/deepgram/v1/listen", + accept_websocket=False, + ) + return websocket, success_handler + + +@pytest.mark.asyncio +async def test_websocket_passthrough_relays_deepgram_transcript_frames_verbatim_and_keeps_them_for_billing(): + """Interim, final and Metadata frames reach the client byte for byte (no JSON round trip, non-ASCII intact, + a binary frame first) and every JSON object frame is what the success handler gets to bill from.""" + upstream_ws = FakeUpstreamWebSocket( + b"\x00\x01binary-first", + DEEPGRAM_INTERIM_FRAME, + "not json at all", + DEEPGRAM_FINAL_FRAME, + DEEPGRAM_METADATA_FRAME, + ) + + websocket, success_handler = await _relay_deepgram_listen(upstream_ws, _pending_receive) + + assert [call.args[0] for call in websocket.send_bytes.await_args_list] == [b"\x00\x01binary-first"] + assert [call.args[0] for call in websocket.send_text.await_args_list] == [ + DEEPGRAM_INTERIM_FRAME, + "not json at all", + DEEPGRAM_FINAL_FRAME, + DEEPGRAM_METADATA_FRAME, + ] + success_call = success_handler.call_args.kwargs + assert success_call["url_route"] == "/deepgram/v1/listen" + assert success_call["response_body"] == [ + json.loads(DEEPGRAM_INTERIM_FRAME), + json.loads(DEEPGRAM_FINAL_FRAME), + json.loads(DEEPGRAM_METADATA_FRAME), + ] + assert success_call["httpx_response"].request.url == DEEPGRAM_LISTEN_TARGET + assert success_call["logging_obj"].model_call_details.get("custom_llm_provider") is None + websocket.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text_upstream_unchanged(): + upstream_ws = RecordingUpstreamWebSocket() + audio_chunk = bytes(range(256)) * 4 + close_stream = json.dumps({"type": "CloseStream"}) + + await _relay_deepgram_listen( + upstream_ws, + AsyncMock( + side_effect=[ + {"type": "websocket.receive", "bytes": audio_chunk}, + {"type": "websocket.receive", "text": close_stream}, + {"type": "websocket.disconnect"}, + ] + ), + ) + + assert [call.args[0] for call in upstream_ws.send.await_args_list] == [audio_chunk, close_stream] + assert isinstance(upstream_ws.send.await_args_list[0].args[0], bytes) + upstream_ws.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_vertex_live_setup_ack_names_the_model_but_is_not_billed_as_usage(): + """Vertex Live keeps its special first frame: the setup acknowledgement is forwarded verbatim, read for the + model, and left out of the frames the usage handler sees; later frames are kept as before.""" + setup_ack = json.dumps( + {"setupComplete": {}, "model": "projects/p/locations/global/publishers/google/models/gemini-live-2.5-flash"} + ) + server_content = json.dumps({"serverContent": {"turnComplete": True}, "usageMetadata": {"totalTokenCount": 12}}) + upstream_ws = FakeUpstreamWebSocket(setup_ack, server_content) + websocket = _client_websocket(_pending_receive) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent", + custom_headers={"Authorization": "Bearer token"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + ) + + assert [call.args[0] for call in websocket.send_text.await_args_list] == [setup_ack, server_content] + success_call = success_handler.call_args.kwargs + assert success_call["response_body"] == [json.loads(server_content)] + assert success_call["logging_obj"].model == "gemini-live-2.5-flash" + assert success_call["logging_obj"].model_call_details["custom_llm_provider"] == "vertex_ai_language_models" + + def _passthrough_kwargs_for_reservation( user_api_key_dict: UserAPIKeyAuth, parsed_body: Optional[dict] = None, From 585c32d3f500d3788cf9a47928bc5922351d835b Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 03:53:28 +0000 Subject: [PATCH 030/102] refactor(deepgram): move listen frame parsing into llms/deepgram and drop routine docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/deepgram/common_utils.py | 63 +++++++++- .../llm_passthrough_endpoints.py | 8 -- ...gram_listen_passthrough_logging_handler.py | 71 +---------- .../pass_through_endpoints.py | 8 -- .../deepgram/test_deepgram_common_utils.py | 114 ++++++++++++++++++ ...gram_listen_passthrough_logging_handler.py | 51 -------- 6 files changed, 179 insertions(+), 136 deletions(-) create mode 100644 tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index db00d048f01..f1759f94775 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -1,5 +1,8 @@ +import math +from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final +from urllib.parse import parse_qs, urlparse import httpx @@ -14,10 +17,6 @@ class DeepgramException(BaseLLMException): def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str: - """ - The upstream ``/listen`` socket for a streaming transcription, keeping the client's query string as sent - and adding the default model only when the client named none - """ listen_url: Final = httpx.URL(f"{(api_base or DEEPGRAM_DEFAULT_API_BASE).rstrip('/')}/listen") websocket_url: Final = listen_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(listen_url.scheme, listen_url.scheme)) params: Final = httpx.QueryParams(query_string) @@ -25,3 +24,59 @@ def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> query_string if params.get("model") else str(params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL)) ) return f"{websocket_url}?{query}" + + +def deepgram_listen_model(upstream_url: str) -> str: + models: Final = parse_qs(urlparse(upstream_url).query).get("model") + return models[0] if models else DEEPGRAM_LISTEN_DEFAULT_MODEL + + +def _seconds(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) if math.isfinite(value) and value >= 0 else None + + +def _results_frame_end(frame: Mapping[str, object]) -> float | None: + start: Final = _seconds(frame.get("start")) + duration: Final = _seconds(frame.get("duration")) + return None if start is None or duration is None else start + duration + + +def _final_transcript(frame: Mapping[str, object]) -> str | None: + if frame.get("is_final") is not True: + return None + channel: Final = frame.get("channel") + alternatives: Final = channel.get("alternatives") if isinstance(channel, Mapping) else None + first: Final = alternatives[0] if isinstance(alternatives, list) and alternatives else None + transcript: Final = first.get("transcript") if isinstance(first, Mapping) else None + return transcript if isinstance(transcript, str) and transcript else None + + +def deepgram_listen_audio_seconds(websocket_messages: Sequence[Mapping[str, object]]) -> float: + metadata_durations: Final = tuple( + duration + for frame in websocket_messages + if frame.get("type") == "Metadata" + if (duration := _seconds(frame.get("duration"))) is not None + ) + if metadata_durations: + return metadata_durations[-1] + return max( + ( + end + for frame in websocket_messages + if frame.get("type") == "Results" + if (end := _results_frame_end(frame)) is not None + ), + default=0.0, + ) + + +def deepgram_listen_transcript(websocket_messages: Sequence[Mapping[str, object]]) -> str: + return " ".join( + transcript + for frame in websocket_messages + if frame.get("type") == "Results" + if (transcript := _final_transcript(frame)) is not None + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index a5a95c34910..1abbf90cb7a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2613,10 +2613,6 @@ def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists: def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: - """ - The first subprotocol the client offered, echoed back so browsers that carry the LiteLLM key in - ``Sec-WebSocket-Protocol`` complete the handshake - """ requested_subprotocols: Final = tuple( protocol.strip() for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",") @@ -2707,10 +2703,6 @@ async def deepgram_listen_websocket_route( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], ) -> None: - """ - Streaming speech to text through Deepgram's ``/v1/listen`` socket. Audio frames and transcript frames are - relayed unchanged; the call is billed on the audio duration Deepgram reports when the socket closes - """ deepgram_api_key: Final = passthrough_endpoint_router.get_credentials( custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value, region_name=None, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py index 6a574a2e1b9..a5fea7c8020 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py @@ -1,81 +1,22 @@ -""" -Cost tracking for Deepgram's streaming ``/v1/listen`` WebSocket. Deepgram bills the audio it processed, which it -reports as ``duration`` on the closing ``Metadata`` frame; a stream that ends without one is billed on the furthest -``start + duration`` across its ``Results`` frames -""" - -import math from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final -from urllib.parse import parse_qs, urlparse +from urllib.parse import urlparse import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import DEEPGRAM_LISTEN_DEFAULT_MODEL from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.deepgram.common_utils import ( + deepgram_listen_audio_seconds, + deepgram_listen_model, + deepgram_listen_transcript, +) from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.types.utils import TranscriptionResponse DEEPGRAM_LISTEN_ROUTE_SUFFIX: Final = "/listen" -def _seconds(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, (int, float)): - return None - return float(value) if math.isfinite(value) and value >= 0 else None - - -def _results_frame_end(frame: Mapping[str, object]) -> float | None: - start: Final = _seconds(frame.get("start")) - duration: Final = _seconds(frame.get("duration")) - return None if start is None or duration is None else start + duration - - -def _final_transcript(frame: Mapping[str, object]) -> str | None: - if frame.get("is_final") is not True: - return None - channel: Final = frame.get("channel") - alternatives: Final = channel.get("alternatives") if isinstance(channel, Mapping) else None - first: Final = alternatives[0] if isinstance(alternatives, list) and alternatives else None - transcript: Final = first.get("transcript") if isinstance(first, Mapping) else None - return transcript if isinstance(transcript, str) and transcript else None - - -def deepgram_listen_audio_seconds(websocket_messages: Sequence[Mapping[str, object]]) -> float: - metadata_durations: Final = tuple( - duration - for frame in websocket_messages - if frame.get("type") == "Metadata" - if (duration := _seconds(frame.get("duration"))) is not None - ) - if metadata_durations: - return metadata_durations[-1] - return max( - ( - end - for frame in websocket_messages - if frame.get("type") == "Results" - if (end := _results_frame_end(frame)) is not None - ), - default=0.0, - ) - - -def deepgram_listen_transcript(websocket_messages: Sequence[Mapping[str, object]]) -> str: - return " ".join( - transcript - for frame in websocket_messages - if frame.get("type") == "Results" - if (transcript := _final_transcript(frame)) is not None - ) - - -def deepgram_listen_model(upstream_url: str) -> str: - models: Final = parse_qs(urlparse(upstream_url).query).get("model") - return models[0] if models else DEEPGRAM_LISTEN_DEFAULT_MODEL - - def _audio_cost(response: TranscriptionResponse, model: str) -> float | None: try: return litellm.completion_cost( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index cf6985852f2..449ae48b0ef 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2121,9 +2121,6 @@ def _resolved_vertex_live_setup( def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: - """ - The frame as a JSON object when it is one, for cost tracking; audio and non-object frames yield None - """ try: decoded: Final = json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")) except (json.JSONDecodeError, UnicodeDecodeError): @@ -2431,10 +2428,6 @@ async def websocket_passthrough_request( json_frame_ordinal: Final = count() async def relay_upstream_frame(upstream_message: str | bytes) -> None: - """ - Send the frame to the client exactly as received, then keep it for cost tracking when it is a JSON - object; the Vertex AI Live setup acknowledgement only names the model, so it is read instead of kept - """ if isinstance(upstream_message, bytes): await websocket.send_bytes(upstream_message) else: @@ -2448,7 +2441,6 @@ async def websocket_passthrough_request( websocket_messages.append(message_data) async def forward_upstream_to_client() -> Close | None: - """Relay upstream frames to the client until the upstream closes, returning its close frame""" try: while True: await relay_upstream_frame(await upstream_ws.recv()) diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py new file mode 100644 index 00000000000..a86cb83d628 --- /dev/null +++ b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py @@ -0,0 +1,114 @@ +import math +from collections.abc import Mapping, Sequence +from typing import Final + +import pytest + +import litellm +from litellm.llms.deepgram.common_utils import ( + deepgram_listen_audio_seconds, + deepgram_listen_model, + deepgram_listen_transcript, + deepgram_listen_websocket_target, +) + +NOVA_3_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" + + +def _results(start: object, duration: object, transcript: str = "", is_final: object = True) -> dict[str, object]: + return { + "type": "Results", + "start": start, + "duration": duration, + "is_final": is_final, + "channel": {"alternatives": [{"transcript": transcript, "confidence": 0.9}]}, + } + + +def _metadata(duration: object) -> dict[str, object]: + return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": 1} + + +@pytest.mark.parametrize( + ("api_base", "query_string", "expected"), + [ + pytest.param( + None, + "model=nova-3&encoding=linear16", + "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16", + id="default", + ), + pytest.param( + None, + "encoding=linear16&sample_rate=16000", + "wss://api.deepgram.com/v1/listen?encoding=linear16&sample_rate=16000&model=nova-3", + id="model added when missing", + ), + pytest.param( + None, + "model=&encoding=linear16", + "wss://api.deepgram.com/v1/listen?encoding=linear16&model=nova-3", + id="empty model replaced", + ), + pytest.param( + "http://localhost:9000/v1/", + "model=nova-2", + "ws://localhost:9000/v1/listen?model=nova-2", + id="custom base becomes ws", + ), + pytest.param( + "wss://dg.internal/v1", + "model=nova-3&keywords=a&keywords=b", + "wss://dg.internal/v1/listen?model=nova-3&keywords=a&keywords=b", + id="repeated keys preserved", + ), + ], +) +def test_deepgram_listen_websocket_target(api_base: str | None, query_string: str, expected: str): + assert deepgram_listen_websocket_target(api_base=api_base, query_string=query_string) == expected + + +@pytest.mark.parametrize( + ("frames", "expected_seconds"), + [ + pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _metadata(6.25)), 6.25, id="metadata wins"), + pytest.param((_metadata(4.0), _results(0.0, 9.0), _metadata(5.5)), 5.5, id="last metadata wins"), + pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _results(1.0, 1.0)), 5.5, id="furthest results end"), + pytest.param((_results(0.0, 0.0), _metadata(0.0)), 0.0, id="zero metadata is a real zero"), + pytest.param((_results(0.0, 1.5), _metadata("6.25")), 1.5, id="string metadata is ignored"), + pytest.param((_results(0.0, 1.5), _metadata(True)), 1.5, id="boolean metadata is ignored"), + pytest.param((_results(0.0, 1.5), _metadata(-3.0)), 1.5, id="negative metadata is ignored"), + pytest.param((_results(0.0, 1.5), _metadata(math.nan), _metadata(math.inf)), 1.5, id="nan/inf ignored"), + pytest.param((_results("0", 2.0), _results(0.0, None), _results(0.0, 0.75)), 0.75, id="malformed results"), + pytest.param(({"type": "SpeechStarted", "timestamp": 3.0}, {"type": "UtteranceEnd"}), 0.0, id="no usage"), + pytest.param((), 0.0, id="no frames"), + ], +) +def test_deepgram_listen_audio_seconds(frames: Sequence[Mapping[str, object]], expected_seconds: float): + assert deepgram_listen_audio_seconds(frames) == expected_seconds + + +def test_deepgram_listen_transcript_joins_final_results_only(): + frames = ( + _results(0.0, 1.0, "hello wor", is_final=False), + _results(0.0, 1.5, "hello world"), + _results(1.5, 0.5, "", is_final=True), + _results(2.0, 1.0, "how are you", is_final="yes"), + {"type": "Results", "start": 3.0, "duration": 1.0, "is_final": True, "channel": {"alternatives": []}}, + _results(4.0, 1.0, "goodbye"), + _metadata(5.0), + ) + assert deepgram_listen_transcript(frames) == "hello world goodbye" + + +@pytest.mark.parametrize( + ("upstream_url", "expected_model"), + [ + (NOVA_3_URL, "nova-3"), + ("wss://api.deepgram.com/v1/listen?encoding=linear16&model=nova-2-medical", "nova-2-medical"), + ("wss://api.deepgram.com/v1/listen?model=nova-3&model=nova-2", "nova-3"), + ("wss://api.deepgram.com/v1/listen?encoding=linear16", litellm.constants.DEEPGRAM_LISTEN_DEFAULT_MODEL), + ], +) +def test_deepgram_listen_model_comes_from_the_upstream_query(upstream_url: str, expected_model: str): + assert deepgram_listen_model(upstream_url) == expected_model diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py index f00411ac10c..40f520e344d 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py @@ -1,7 +1,5 @@ """Deepgram ``/v1/listen`` WebSocket passthrough: duration extraction and duration based cost tracking.""" -import math -from collections.abc import Mapping, Sequence from datetime import datetime from types import SimpleNamespace from typing import Final @@ -14,9 +12,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_provider_handlers.deepgram_listen_passthrough_logging_handler import ( DeepgramListenPassthroughLoggingHandler, - deepgram_listen_audio_seconds, - deepgram_listen_model, - deepgram_listen_transcript, ) from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging from litellm.types.passthrough_endpoints.pass_through_endpoints import PassthroughStandardLoggingPayload @@ -39,52 +34,6 @@ def _metadata(duration: object) -> dict[str, object]: return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": 1} -@pytest.mark.parametrize( - ("frames", "expected_seconds"), - [ - pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _metadata(6.25)), 6.25, id="metadata wins"), - pytest.param((_metadata(4.0), _results(0.0, 9.0), _metadata(5.5)), 5.5, id="last metadata wins"), - pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _results(1.0, 1.0)), 5.5, id="furthest results end"), - pytest.param((_results(0.0, 0.0), _metadata(0.0)), 0.0, id="zero metadata is a real zero"), - pytest.param((_results(0.0, 1.5), _metadata("6.25")), 1.5, id="string metadata is ignored"), - pytest.param((_results(0.0, 1.5), _metadata(True)), 1.5, id="boolean metadata is ignored"), - pytest.param((_results(0.0, 1.5), _metadata(-3.0)), 1.5, id="negative metadata is ignored"), - pytest.param((_results(0.0, 1.5), _metadata(math.nan), _metadata(math.inf)), 1.5, id="nan/inf ignored"), - pytest.param((_results("0", 2.0), _results(0.0, None), _results(0.0, 0.75)), 0.75, id="malformed results"), - pytest.param(({"type": "SpeechStarted", "timestamp": 3.0}, {"type": "UtteranceEnd"}), 0.0, id="no usage"), - pytest.param((), 0.0, id="no frames"), - ], -) -def test_deepgram_listen_audio_seconds(frames: Sequence[Mapping[str, object]], expected_seconds: float): - assert deepgram_listen_audio_seconds(frames) == expected_seconds - - -def test_deepgram_listen_transcript_joins_final_results_only(): - frames = ( - _results(0.0, 1.0, "hello wor", is_final=False), - _results(0.0, 1.5, "hello world"), - _results(1.5, 0.5, "", is_final=True), - _results(2.0, 1.0, "how are you", is_final="yes"), - {"type": "Results", "start": 3.0, "duration": 1.0, "is_final": True, "channel": {"alternatives": []}}, - _results(4.0, 1.0, "goodbye"), - _metadata(5.0), - ) - assert deepgram_listen_transcript(frames) == "hello world goodbye" - - -@pytest.mark.parametrize( - ("upstream_url", "expected_model"), - [ - (NOVA_3_URL, "nova-3"), - ("wss://api.deepgram.com/v1/listen?encoding=linear16&model=nova-2-medical", "nova-2-medical"), - ("wss://api.deepgram.com/v1/listen?model=nova-3&model=nova-2", "nova-3"), - ("wss://api.deepgram.com/v1/listen?encoding=linear16", litellm.constants.DEEPGRAM_LISTEN_DEFAULT_MODEL), - ], -) -def test_deepgram_listen_model_comes_from_the_upstream_query(upstream_url: str, expected_model: str): - assert deepgram_listen_model(upstream_url) == expected_model - - @pytest.mark.parametrize( ("url_route", "expected"), [ From 9fa85c5da5a6a9497a0df12cbc02866497af353d Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 08:39:01 +0000 Subject: [PATCH 031/102] fix(proxy): name the blocking guardrail in x-litellm-applied-guardrails When a guardrail hook raises, the common ProxyLogging dispatch (sequential and parallel pre_call, pipeline block, during_call and post_call metrics wrapper, streaming iterator wrapper) now records that guardrail in applied_guardrails before re-raising, and pre_call_hook folds request-declared guardrails in on its raising path. Buffered streams rebuild their response headers after the first chunk so a post_call block reached while buffering carries the blocker too Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 11 ++-- litellm/proxy/utils.py | 58 +++++++++++++++---- .../proxy/test_common_request_processing.py | 57 ++++++++++++++++++ .../proxy_logging/test_during_call_hook.py | 18 ++++++ .../proxy_logging/test_guardrail_pipeline.py | 8 ++- .../test_post_call_success_hook.py | 24 ++++++++ .../utils/proxy_logging/test_pre_call_hook.py | 40 +++++++++++++ .../proxy_logging/test_streaming_hooks.py | 38 +++++++++++- 8 files changed, 234 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2f39e6c71bc..9a038ca79ba 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2576,11 +2576,14 @@ class ProxyBaseLLMRequestProcessing: ) async def refresh_stream_headers() -> Mapping[str, str]: - """`custom_headers` rebuilt for whichever deployment served the stream.""" - if not getattr(response, "fallback_headers_adopted", False): - return custom_headers + """`custom_headers` rebuilt once the first chunk is buffered, from `self.data` as the + guardrails left it and for whichever deployment served the stream.""" return self._stream_response_headers( - hidden_params=get_hidden_params_dict(response), + hidden_params=( + get_hidden_params_dict(response) + if getattr(response, "fallback_headers_adopted", False) + else hidden_params + ), user_api_key_dict=user_api_key_dict, logging_obj=logging_obj, version=version, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8225fef3492..fea6ce20b61 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -123,6 +123,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( @@ -437,6 +438,12 @@ def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: detail.setdefault("guardrail_mode", event_hook) +def _record_raising_guardrail(request_data: Mapping[str, object], callback: object) -> None: + guardrail_name: Final[object] = getattr(callback, "guardrail_name", None) + if isinstance(request_data, dict) and isinstance(guardrail_name, str): + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=guardrail_name) + + def _is_client_error_exception(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code < 500 @@ -1795,13 +1802,19 @@ class ProxyLogging: ) if expected_if_unmutated is not None: callback.mark_pre_call_hook_ran(expected_if_unmutated) - result: Final = await self._process_guardrail_callback( - callback=callback, - data=input_data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - event_type=GuardrailEventHooks.pre_call, - ) + try: + result: Final = await self._process_guardrail_callback( + callback=callback, + data=input_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + event_type=GuardrailEventHooks.pre_call, + ) + except SensitiveDataRouteException: + raise + except Exception: + _record_raising_guardrail(data, callback) + raise if ( scans_raw_request and expected_if_unmutated is not None @@ -2031,6 +2044,7 @@ class ProxyLogging: callback: Final = PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) if callback is not None: _enrich_http_exception_with_guardrail_context(original_exception, callback) + _record_raising_guardrail(data, callback) raise original_exception step_results_serializable: Final = [ @@ -2296,8 +2310,10 @@ class ProxyLogging: if data is not None: self._process_guardrail_metadata(data) return data - except Exception as e: - raise e + except Exception: + if data is not None: + self._process_guardrail_metadata(data) + raise async def _run_parallel_pre_call_guardrails( self, @@ -2355,6 +2371,8 @@ class ProxyLogging: # live kwargs. if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None: callback.mark_pre_call_hook_ran(data) + if isinstance(result, BaseException) and not isinstance(result, SensitiveDataRouteException): + _record_raising_guardrail(data, callback) raised: Final = tuple(result for result in results if isinstance(result, BaseException)) blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) if blocking is not None: @@ -2433,7 +2451,12 @@ class ProxyLogging: break @staticmethod - async def _run_guardrail_with_metrics(callback: object, coro: Awaitable[_T], hook_type: str) -> _T: + async def _run_guardrail_with_metrics( + callback: object, + coro: Awaitable[_T], + hook_type: str, + request_data: Mapping[str, object], + ) -> _T: """ Await `coro`, recording its latency and status to the `litellm_guardrail_latency_seconds` metric under `hook_type`, and @@ -2453,6 +2476,7 @@ class ProxyLogging: status = "error" error_type = type(e).__name__ _enrich_http_exception_with_guardrail_context(e, callback) + _record_raising_guardrail(request_data, callback) raise finally: ProxyLogging._emit_guardrail_metrics( @@ -2465,7 +2489,9 @@ class ProxyLogging: @staticmethod async def _wrap_streaming_iterator_with_enrichment( - callback: object, gen: AsyncGenerator[_T, None] + callback: object, + gen: AsyncGenerator[_T, None], + request_data: Mapping[str, object], ) -> AsyncGenerator[_T, None]: """ Yield from `gen`; if iteration raises an HTTPException with dict detail, @@ -2480,6 +2506,7 @@ class ProxyLogging: yield chunk except Exception as e: _enrich_http_exception_with_guardrail_context(e, callback) + _record_raising_guardrail(request_data, callback) raise # Cache for callback-capability detection. Keyed on a signature of @@ -2714,6 +2741,7 @@ class ProxyLogging: call_type=call_type, ), "during_call", + request_data=data, ) return await self._run_guardrail_with_metrics( @@ -2724,6 +2752,7 @@ class ProxyLogging: call_type=call_type, ), "during_call", + request_data=data, ) async def failed_tracking_alert( @@ -3242,6 +3271,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) else: guardrail_response = await self._run_guardrail_with_metrics( @@ -3252,6 +3282,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) if guardrail_response is not None: @@ -3315,6 +3346,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) else: await self._run_guardrail_with_metrics( @@ -3325,6 +3357,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) results: Final = await asyncio.gather( @@ -3388,6 +3421,7 @@ class ProxyLogging: request_data=request_data, ), "post_mcp_call", + request_data=request_data, ) return response @@ -3637,6 +3671,7 @@ class ProxyLogging: response=current_response, request_data=request_data, ), + request_data=request_data, ) else: # kind == "apply_guardrail": route through unified_guardrail @@ -3649,6 +3684,7 @@ class ProxyLogging: guardrail_to_apply=resolved_callback, buffer_until_moderated_default=(kind == "override"), ), + request_data=request_data, ) pipeline_translation: Final = ( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 4ac687625c2..028ea29c093 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -40,9 +40,12 @@ from litellm.proxy.common_request_processing import ( _parse_event_data_for_error, _resolve_per_request_model_group_alias, _should_return_raw_model_name, + _sse_error_frames, _UpstreamClosingStreamingResponse, create_response, + sse_error_payload, ) +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -8952,6 +8955,60 @@ class TestStreamingResponseHeadersFollowFallback: assert "llm_provider-stale-marker" not in result.headers assert result.headers["x-callback-header"] == "kept" + @pytest.mark.asyncio + async def test_streaming_block_headers_name_the_blocking_guardrail(self, monkeypatch): + processor_data: dict[str, object] = {"model": "oa", "stream": True, "metadata": {}} + + def select_data_generator(**kwargs): + async def generator(): + add_guardrail_to_applied_guardrails_header(processor_data, "stream-blocker") + _, error_obj = sse_error_payload(HTTPException(status_code=400, detail="blocked")) + for frame in _sse_error_frames(error_obj): + yield frame + + return generator() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "lit-7144-call" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + processor_data["litellm_logging_obj"] = logging_obj + processor = ProxyBaseLLMRequestProcessing(data=processor_data) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + async def fake_route_request(**kwargs): + async def call(): + return SimpleNamespace(_hidden_params={}, fallback_headers_adopted=False) + + return call() + + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + + result = await processor.base_process_llm_request( + request=Request(scope={"type": "http", "headers": []}), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + is_streaming_request=True, + skip_pre_call_logic=True, + ) + + assert isinstance(result, JSONResponse) + assert result.status_code == 400 + assert result.headers["x-litellm-applied-guardrails"] == "stream-blocker" + class _MessagesFallbackStream: def __init__(self) -> None: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py index 3c5d879c2dc..46f39ef6fb7 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py @@ -6,6 +6,7 @@ from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -84,3 +85,20 @@ async def test_during_call_hook_guardrail_error_raises(proxy_logging, make_user_ user_api_key_dict=make_user_api_key_auth(), call_type="completion", ) + + +@pytest.mark.asyncio +async def test_during_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + g = _make_guardrail("blocker") + g.async_moderation_hook = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked")) + monkeypatch.setattr(litellm, "callbacks", [_make_guardrail("passer"), g]) + data = {"model": "m", "metadata": {}} + with pytest.raises(HTTPException): + await proxy_logging.during_call_hook( + data=data, + user_api_key_dict=make_user_api_key_auth(), + call_type="completion", + ) + assert "blocker" in data["metadata"]["applied_guardrails"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 077bf5a313e..cd2b7a278bb 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -527,17 +527,19 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode(): result.step_results = [MagicMock(guardrail_name="g")] result.original_exception = original + data: dict[str, object] = {"model": "m"} saved = litellm.callbacks litellm.callbacks = [cb] try: with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") finally: litellm.callbacks = saved assert info.value is original assert info.value.detail["guardrail_name"] == "g" assert info.value.detail["guardrail_mode"] == GuardrailEventHooks.pre_call + assert data["metadata"] == {"applied_guardrails": ["g"]} def test_handle_pipeline_result_block_does_not_reraise_sensitive_data_route(): @@ -617,7 +619,7 @@ async def test_run_guardrail_with_metrics_passes_result_and_records_success(monk monkeypatch.setattr(litellm, "callbacks", [prom]) out = await ProxyLogging._run_guardrail_with_metrics( - callback=MagicMock(guardrail_name="g"), coro=task(), hook_type="during_call" + callback=MagicMock(guardrail_name="g"), coro=task(), hook_type="during_call", request_data={} ) assert out == {"a": 1, "b": 2, "c": 3} @@ -643,7 +645,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch monkeypatch.setattr(litellm, "callbacks", [prom]) with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call") + await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call", request_data={}) assert detail["guardrail_name"] == "presidio" recorded = prom._record_guardrail_metrics.call_args.kwargs diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py index 715d66db181..53d8948869f 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -6,6 +6,7 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -96,3 +97,26 @@ async def test_post_call_success_hook_guardrail_returns_modified_response( data={}, response={"orig": True}, user_api_key_dict=make_user_api_key_auth() ) assert out == modified + + +@pytest.mark.asyncio +@pytest.mark.parametrize("run_in_parallel", [False, True], ids=["sequential", "parallel"]) +async def test_post_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch, run_in_parallel +): + def _passer_that_records(data, user_api_key_dict, response): + data["metadata"]["applied_guardrails"] = ["passer"] + + passer = _make_guardrail("passer") + passer.async_post_call_success_hook = AsyncMock(side_effect=_passer_that_records) + passer.run_in_parallel = run_in_parallel + blocker = _make_guardrail("blocker") + blocker.async_post_call_success_hook = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked")) + blocker.run_in_parallel = run_in_parallel + monkeypatch.setattr(litellm, "callbacks", [passer, blocker]) + data = {"model": "m", "metadata": {}} + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data=data, response=MagicMock(), user_api_key_dict=make_user_api_key_auth() + ) + assert data["metadata"]["applied_guardrails"] == ["passer", "blocker"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 6e5cb7fcae3..dbc6fba4ab1 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -905,3 +905,43 @@ async def test_scan_raw_request_warns_on_in_place_mutation_returning_none( ) mock_logger.warning.assert_called_once() assert "scan_raw_request" in str(mock_logger.warning.call_args) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "blocker_kwargs", + [ + pytest.param({}, id="sequential"), + pytest.param({"scan_raw_request": True}, id="scan_raw_request"), + pytest.param({"run_in_parallel": True}, id="parallel"), + ], +) +async def test_pre_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch, blocker_kwargs +): + monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(**blocker_kwargs)]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + data = _secret_request() + with pytest.raises(HTTPException, match="blocked"): + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + ) + assert data["metadata"]["applied_guardrails"] == ["blocker"] + + +@pytest.mark.asyncio +async def test_pre_call_block_keeps_request_declared_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(default_on=False)]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + data = {**_secret_request(), "metadata": {"guardrails": ["blocker", "declared-post-call"]}} + with pytest.raises(HTTPException, match="blocked"): + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + ) + assert data["metadata"]["applied_guardrails"] == ["blocker", "declared-post-call"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index ebc831b4102..52586ed2174 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -20,6 +20,7 @@ from fastapi import HTTPException import litellm from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( @@ -27,6 +28,7 @@ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterato ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import Usage @@ -175,7 +177,7 @@ async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(pro yield ch cb = MagicMock(guardrail_name="g", event_hook="pre_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=gen()) + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=gen(), request_data={}) out = [ch async for ch in wrapped] snapshot = { "chunks": out, @@ -201,7 +203,7 @@ async def test_wrap_streaming_iterator_with_enrichment_enriches_http_exception_r raise HTTPException(status_code=400, detail=detail) cb = MagicMock(guardrail_name="presidio", event_hook="post_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=boom_gen()) + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=boom_gen(), request_data={}) with pytest.raises(HTTPException): async for _ in wrapped: pass @@ -696,3 +698,35 @@ async def test_post_call_response_headers_hook_swallows_callback_error(proxy_log data={}, user_api_key_dict=make_user_api_key_auth(), response=response ) assert out == {} + + +@pytest.mark.asyncio +async def test_stream_guardrail_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class _StreamBlocker(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="stream-blocker", event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for _ in response: + raise HTTPException(status_code=400, detail={"error": "blocked"}) + yield # pragma: no cover + + monkeypatch.setattr(litellm, "callbacks", [_StreamBlocker()]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + async def upstream(): + yield "chunk" + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(HTTPException): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream(), + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + assert request_data["metadata"]["applied_guardrails"] == ["stream-blocker"] From 849859001f5e0cce29bd973778fca71e3350e16b Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 18:18:40 +0000 Subject: [PATCH 032/102] fix(deepgram): refuse callback delivery on the /listen passthrough so sessions cannot go unbilled With callback or callback_method in the query, Deepgram sends every Results and Metadata frame to the caller's URL and only a request id down this socket, so the proxy would meter zero seconds of audio while its own Deepgram credential paid for the transcription. The route now closes such connections with 1008 before contacting Deepgram, naming the offending parameters in the close reason. Adds helper and route tests for both parameters and a nine mutation sweep, all killed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/deepgram/common_utils.py | 5 +++ .../llm_passthrough_endpoints.py | 14 +++++- .../deepgram/test_deepgram_common_utils.py | 19 ++++++++ .../test_deepgram_ws_passthrough_routes.py | 45 +++++++++++++++++++ 4 files changed, 82 insertions(+), 1 deletion(-) diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index f1759f94775..947df37bbe7 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -10,6 +10,7 @@ from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT from litellm.llms.base_llm.chat.transformation import BaseLLMException _WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"}) +DEEPGRAM_LISTEN_CALLBACK_PARAMS: Final = frozenset({"callback", "callback_method"}) class DeepgramException(BaseLLMException): @@ -26,6 +27,10 @@ def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> return f"{websocket_url}?{query}" +def deepgram_listen_callback_params(query_string: str) -> tuple[str, ...]: + return tuple(sorted(DEEPGRAM_LISTEN_CALLBACK_PARAMS.intersection(httpx.QueryParams(query_string).keys()))) + + def deepgram_listen_model(upstream_url: str) -> str: models: Final = parse_qs(urlparse(upstream_url).query).get("model") return models[0] if models else DEEPGRAM_LISTEN_DEFAULT_MODEL diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1abbf90cb7a..7f9e0169fb2 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -36,7 +36,10 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.llms.deepgram.common_utils import deepgram_listen_websocket_target +from litellm.llms.deepgram.common_utils import ( + deepgram_listen_callback_params, + deepgram_listen_websocket_target, +) from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse @@ -2694,6 +2697,7 @@ async def openai_websocket_proxy_route( _DEEPGRAM_WS_MISSING_KEY_REASON: Final = ( "Required 'DEEPGRAM_API_KEY' in environment to make pass-through calls to Deepgram." ) +_DEEPGRAM_WS_CALLBACK_REASON: Final = "Deepgram callback delivery is not supported through the proxy: remove {params}" @router.websocket("/deepgram/v1/listen") @@ -2712,6 +2716,14 @@ async def deepgram_listen_websocket_route( return await websocket.accept(subprotocol=_negotiated_websocket_subprotocol(websocket)) + callback_params: Final = deepgram_listen_callback_params(websocket.url.query) + if callback_params: + await websocket.close( + code=1008, + reason=_DEEPGRAM_WS_CALLBACK_REASON.format(params=", ".join(callback_params)), + ) + return + await relay( websocket=websocket, target=deepgram_listen_websocket_target( diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py index a86cb83d628..65fbf7c7870 100644 --- a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py @@ -7,6 +7,7 @@ import pytest import litellm from litellm.llms.deepgram.common_utils import ( deepgram_listen_audio_seconds, + deepgram_listen_callback_params, deepgram_listen_model, deepgram_listen_transcript, deepgram_listen_websocket_target, @@ -68,6 +69,24 @@ def test_deepgram_listen_websocket_target(api_base: str | None, query_string: st assert deepgram_listen_websocket_target(api_base=api_base, query_string=query_string) == expected +@pytest.mark.parametrize( + ("query_string", "expected"), + [ + pytest.param("model=nova-3&encoding=linear16", (), id="no callback"), + pytest.param("model=nova-3&callback=https%3A%2F%2Fevil.example%2Fsink", ("callback",), id="callback"), + pytest.param( + "callback_method=put&model=nova-3&callback=wss%3A%2F%2Fevil.example", + ("callback", "callback_method"), + id="callback and method", + ), + pytest.param("model=nova-3&callback_method=put", ("callback_method",), id="method alone"), + pytest.param("model=nova-3&callbacks=x&my_callback=y", (), id="only exact names match"), + ], +) +def test_deepgram_listen_callback_params(query_string: str, expected: tuple[str, ...]): + assert deepgram_listen_callback_params(query_string) == expected + + @pytest.mark.parametrize( ("frames", "expected_seconds"), [ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 64804e621ba..5ea2b0b8ab9 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -207,6 +207,30 @@ async def test_deepgram_listen_closes_cleanly_when_provider_credentials_missing( assert relay.calls == [] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "query", + [ + pytest.param("model=nova-3&callback=https%3A%2F%2Fsink.example%2Fdg", id="http callback"), + pytest.param("callback=wss%3A%2F%2Fsink.example&callback_method=put&model=nova-3", id="ws callback"), + ], +) +async def test_deepgram_listen_rejects_callback_delivery_that_would_go_unbilled(query, monkeypatch): + """With ``callback`` set, Deepgram sends every Results and Metadata frame to the caller's URL and only a + request id down this socket, so the proxy would meter zero seconds of audio; refuse before contacting Deepgram.""" + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + websocket = _FakeWebSocket("/deepgram/v1/listen", query) + + with patch(GET_CREDENTIALS, return_value="dg-provider-key"): + relay = await _serve(websocket) + + assert relay.calls == [] + assert websocket.closed is not None + assert websocket.closed[0] == 1008 + assert "callback" in websocket.closed[1] + assert "dg-provider-key" not in websocket.closed[1] + + def _app_with_relay(relay: _FakeRelay) -> FastAPI: app = FastAPI() app.include_router(router) @@ -228,6 +252,27 @@ def test_deepgram_listen_rejects_connections_without_a_litellm_key(): get_credentials.assert_not_called() +def test_deepgram_listen_callback_rejection_reaches_the_client_as_a_policy_close(monkeypatch): + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + relay = _FakeRelay() + client = TestClient(_app_with_relay(relay)) + + with ( + patch(GET_CREDENTIALS, return_value="dg-provider-key"), + patch(USER_API_KEY_AUTH, new=AsyncMock(return_value=UserAPIKeyAuth(api_key="hashed"))), + ): + with pytest.raises(WebSocketDisconnect) as disconnect: + with client.websocket_connect( + "/deepgram/v1/listen?model=nova-3&callback=https%3A%2F%2Fsink.example%2Fdg", + headers={"Authorization": "Bearer sk-litellm-virtual"}, + ) as connection: + connection.receive_text() + + assert disconnect.value.code == 1008 + assert "callback" in disconnect.value.reason + assert relay.calls == [] + + def test_deepgram_listen_authenticates_the_litellm_key_and_relays_to_deepgram(monkeypatch): monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) relay = _FakeRelay() From 73dea4c5676218005a42bdbe841c76ce6eb42707 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 18:58:21 +0000 Subject: [PATCH 033/102] fix(proxy): attribute only the raising layer in stream and pipeline blocks The streaming wrapper caught every exception crossing its boundary and named its own callback, so a block by an inner guardrail or a provider stream failure also named every outer guardrail. The wrapper now runs the hook over an upstream boundary that remembers the exception it raised, and skips attribution when the same exception passes through Pipeline blocks converted from SensitiveDataRouteException or ModifyResponseException into a generic guardrail_pipeline_error now still record the blocking step's guardrail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 129 ++++++++++++----- .../proxy_logging/test_guardrail_pipeline.py | 26 +++- .../proxy_logging/test_streaming_hooks.py | 132 +++++++++++++++--- 3 files changed, 223 insertions(+), 64 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index fea6ce20b61..21e6ed1d24f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -12,13 +12,36 @@ import sys import threading import time import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Mapping, + Sequence, +) from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText +from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, TypeVar, Union, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Final, + Generic, + Literal, + Optional, + Protocol, + TypeVar, + Union, + cast, + overload, +) from typing_extensions import ReadOnly, TypedDict @@ -444,6 +467,33 @@ def _record_raising_guardrail(request_data: Mapping[str, object], callback: obje add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=guardrail_name) +class _UpstreamStreamBoundary(Generic[_T]): + """Remembers the exception the upstream iterator raised, so the wrapper around a + streaming hook can tell a pass-through failure from one the hook raised itself.""" + + __slots__ = ("_upstream", "failure") + + def __init__(self, upstream: AsyncIterable[_T]) -> None: + self._upstream: Final = upstream.__aiter__() + self.failure: BaseException | None = None + + def __aiter__(self) -> "_UpstreamStreamBoundary[_T]": + return self + + async def __anext__(self) -> _T: + try: + return await self._upstream.__anext__() + except StopAsyncIteration: + raise + except Exception as e: + self.failure = e + raise + + +class _StreamIteratorHook(Protocol[_T]): + def __call__(self, *, response: AsyncIterator[_T]) -> AsyncGenerator[_T, None]: ... + + def _is_client_error_exception(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code < 500 @@ -2037,14 +2087,18 @@ class ProxyLogging: _merge_pipeline_metadata_writes(data, result.modified_data) if result.terminal_action == "block": + blocking_step: Final = result.step_results[-1] if result.step_results else None + callback: Final = ( + PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) + if blocking_step is not None + else None + ) + if callback is not None: + _record_raising_guardrail(data, callback) original_exception: Final = result.original_exception if original_exception is not None and not _exception_changes_request_flow(original_exception): - blocking_step: Final = result.step_results[-1] if result.step_results else None - if blocking_step is not None: - callback: Final = PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) - if callback is not None: - _enrich_http_exception_with_guardrail_context(original_exception, callback) - _record_raising_guardrail(data, callback) + if callback is not None: + _enrich_http_exception_with_guardrail_context(original_exception, callback) raise original_exception step_results_serializable: Final = [ @@ -2490,23 +2544,26 @@ class ProxyLogging: @staticmethod async def _wrap_streaming_iterator_with_enrichment( callback: object, - gen: AsyncGenerator[_T, None], + response: AsyncIterable[_T], + hook: _StreamIteratorHook[_T], request_data: Mapping[str, object], ) -> AsyncGenerator[_T, None]: """ - Yield from `gen`; if iteration raises an HTTPException with dict detail, - enrich the detail with the originating callback's `guardrail_name` and - `guardrail_mode` before re-raising. Used to wrap each layer of the - async_post_call_streaming_iterator_hook chain so the enrichment is - attributed to the callback that produced the chunk pipeline at that - point in the chain. + Run `hook` over `response` and yield its chunks. If the hook itself raises, + enrich an HTTPException's dict detail with the callback's `guardrail_name` + and `guardrail_mode` and record the callback in `applied_guardrails` before + re-raising. Failures raised by `response` (the provider stream or an inner + layer of the async_post_call_streaming_iterator_hook chain) pass through + untouched, so only the layer that actually raised is attributed. """ + upstream: Final = _UpstreamStreamBoundary(response) try: - async for chunk in gen: + async for chunk in hook(response=upstream): yield chunk except Exception as e: - _enrich_http_exception_with_guardrail_context(e, callback) - _record_raising_guardrail(request_data, callback) + if e is not upstream.failure: + _enrich_http_exception_with_guardrail_context(e, callback) + _record_raising_guardrail(request_data, callback) raise # Cache for callback-capability detection. Keyed on a signature of @@ -3663,29 +3720,27 @@ class ProxyLogging: ) else kind ) - if effective_kind == "override": - current_response = self._wrap_streaming_iterator_with_enrichment( - resolved_callback, - resolved_callback.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=current_response, - request_data=request_data, - ), + hook: _StreamIteratorHook[object] = ( + partial( + resolved_callback.async_post_call_streaming_iterator_hook, + user_api_key_dict=user_api_key_dict, request_data=request_data, ) - else: - # kind == "apply_guardrail": route through unified_guardrail - current_response = self._wrap_streaming_iterator_with_enrichment( - resolved_callback, - unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - request_data=request_data, - response=current_response, - guardrail_to_apply=resolved_callback, - buffer_until_moderated_default=(kind == "override"), - ), + if effective_kind == "override" + else partial( + unified_guardrail.async_post_call_streaming_iterator_hook, + user_api_key_dict=user_api_key_dict, request_data=request_data, + guardrail_to_apply=resolved_callback, + buffer_until_moderated_default=(kind == "override"), ) + ) + current_response = self._wrap_streaming_iterator_with_enrichment( + resolved_callback, + current_response, + hook, + request_data=request_data, + ) pipeline_translation: Final = ( resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index cd2b7a278bb..5f3c09d9195 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -551,14 +551,23 @@ def test_handle_pipeline_result_block_does_not_reraise_sensitive_data_route(): session_id="sess-1", guardrail_name="pii-router", ) + cb = _make_guardrail() + cb.guardrail_name = "pii-router" result = MagicMock() result.terminal_action = "block" result.step_results = [MagicMock(guardrail_name="pii-router")] result.original_exception = original - with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + data: dict[str, object] = {"model": "m"} + saved = litellm.callbacks + litellm.callbacks = [cb] + try: + with pytest.raises(HTTPException) as info: + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") + finally: + litellm.callbacks = saved assert info.value.status_code == 400 assert info.value.detail["error"]["type"] == "guardrail_pipeline_error" + assert data["metadata"] == {"applied_guardrails": ["pii-router"]} def test_handle_pipeline_result_block_does_not_reraise_modify_response(): @@ -571,14 +580,23 @@ def test_handle_pipeline_result_block_does_not_reraise_modify_response(): request_data={"model": "m"}, guardrail_name="masker", ) + cb = _make_guardrail() + cb.guardrail_name = "masker" result = MagicMock() result.terminal_action = "block" result.step_results = [MagicMock(guardrail_name="masker")] result.original_exception = original - with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + data: dict[str, object] = {"model": "m"} + saved = litellm.callbacks + litellm.callbacks = [cb] + try: + with pytest.raises(HTTPException) as info: + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") + finally: + litellm.callbacks = saved assert info.value.status_code == 400 assert info.value.detail["error"]["type"] == "guardrail_pipeline_error" + assert data["metadata"] == {"applied_guardrails": ["masker"]} def test_handle_pipeline_result_modify_response_raises_modify_exception(): diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index 52586ed2174..6fb000b4fa7 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -170,6 +170,15 @@ def test_init_response_taking_too_long_task_no_slack_instance_no_error_raises(pr # --------------------------------------------------------------------------- +async def _passthrough_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]: + async for chunk in response: + yield chunk + + +async def _one_chunk() -> AsyncGenerator[object, None]: + yield "chunk" + + @pytest.mark.asyncio async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(proxy_logging): async def gen(): @@ -177,7 +186,9 @@ async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(pro yield ch cb = MagicMock(guardrail_name="g", event_hook="pre_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=gen(), request_data={}) + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=gen(), hook=_passthrough_hook, request_data={} + ) out = [ch async for ch in wrapped] snapshot = { "chunks": out, @@ -197,18 +208,43 @@ async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(pro async def test_wrap_streaming_iterator_with_enrichment_enriches_http_exception_raises(proxy_logging): detail = {"error": "blocked"} - async def boom_gen(): + async def boom_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]: if False: yield # pragma: no cover raise HTTPException(status_code=400, detail=detail) cb = MagicMock(guardrail_name="presidio", event_hook="post_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=boom_gen(), request_data={}) + request_data: dict[str, object] = {} + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=_one_chunk(), hook=boom_hook, request_data=request_data + ) with pytest.raises(HTTPException): async for _ in wrapped: pass assert detail["guardrail_name"] == "presidio" assert detail["guardrail_mode"] == "post_call" + assert request_data["metadata"]["applied_guardrails"] == ["presidio"] + + +@pytest.mark.asyncio +async def test_wrap_streaming_iterator_leaves_upstream_http_exception_unattributed(proxy_logging): + detail = {"error": "upstream rejected the stream"} + + async def failing_upstream() -> AsyncGenerator[object, None]: + if False: + yield # pragma: no cover + raise HTTPException(status_code=502, detail=detail) + + cb = MagicMock(guardrail_name="presidio", event_hook="post_call") + request_data: dict[str, object] = {} + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=failing_upstream(), hook=_passthrough_hook, request_data=request_data + ) + with pytest.raises(HTTPException): + async for _ in wrapped: + pass + assert detail == {"error": "upstream rejected the stream"} + assert request_data == {} # --------------------------------------------------------------------------- @@ -700,33 +736,83 @@ async def test_post_call_response_headers_hook_swallows_callback_error(proxy_log assert out == {} +class _StreamBlocker(CustomGuardrail): + def __init__(self, guardrail_name: str = "stream-blocker") -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for _ in response: + raise HTTPException(status_code=400, detail={"error": "blocked"}) + yield # pragma: no cover + + +class _StreamPasser(CustomGuardrail): + def __init__(self, guardrail_name: str = "stream-passer") -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for chunk in response: + yield chunk + + +async def _drain_stream_chain( + proxy_logging: ProxyLogging, + user_api_key_dict: UserAPIKeyAuth, + upstream: AsyncIterator[object], + request_data: dict[str, object], +) -> None: + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ): + pass + + +async def _failing_provider_stream() -> AsyncGenerator[object, None]: + yield "chunk" + raise RuntimeError("provider connection dropped") + + @pytest.mark.asyncio async def test_stream_guardrail_block_names_the_blocking_guardrail_in_applied_guardrails( proxy_logging, make_user_api_key_auth, monkeypatch ): - class _StreamBlocker(CustomGuardrail): - def __init__(self) -> None: - super().__init__(guardrail_name="stream-blocker", event_hook=GuardrailEventHooks.post_call, default_on=True) - - async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] - ) -> AsyncGenerator[object, None]: - async for _ in response: - raise HTTPException(status_code=400, detail={"error": "blocked"}) - yield # pragma: no cover - monkeypatch.setattr(litellm, "callbacks", [_StreamBlocker()]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) - async def upstream(): - yield "chunk" - request_data: dict[str, object] = {"metadata": {}} with pytest.raises(HTTPException): - async for _ in proxy_logging.async_post_call_streaming_iterator_hook( - response=upstream(), - user_api_key_dict=make_user_api_key_auth(), - request_data=request_data, - ): - pass + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _one_chunk(), request_data) assert request_data["metadata"]["applied_guardrails"] == ["stream-blocker"] + + +@pytest.mark.asyncio +async def test_stream_block_by_inner_guardrail_does_not_name_the_outer_layers( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_StreamBlocker(), _StreamPasser("outer-a"), _StreamPasser("outer-b")]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(HTTPException) as info: + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _one_chunk(), request_data) + assert info.value.detail["guardrail_name"] == "stream-blocker" + assert request_data["metadata"]["applied_guardrails"] == ["stream-blocker"] + + +@pytest.mark.asyncio +async def test_stream_provider_failure_is_not_attributed_to_any_guardrail( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_StreamPasser("outer-a"), _StreamPasser("outer-b")]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(RuntimeError, match="provider connection dropped"): + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _failing_provider_stream(), request_data) + assert request_data["metadata"] == {} From db560ca65248d4c03c0c02a1d4cca392d0294144 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 19:02:44 +0000 Subject: [PATCH 034/102] fix(passthrough): bill every Deepgram channel, not just wall-clock duration Deepgram charges for the total processed audio across channels, so a stereo /listen session with multichannel=true costs twice its duration. The handler now multiplies the session duration by a validated channel count taken from Metadata.channels, then the widest Results channel_index, then the channels query parameter, defaulting to one. Booleans, floats, strings, zero and negative values are ignored Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/deepgram/common_utils.py | 40 +++++++++++ ...gram_listen_passthrough_logging_handler.py | 8 ++- .../deepgram/test_deepgram_common_utils.py | 70 ++++++++++++++++++- ...gram_listen_passthrough_logging_handler.py | 32 ++++++++- 4 files changed, 143 insertions(+), 7 deletions(-) diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index 947df37bbe7..6beb70b4499 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -36,6 +36,46 @@ def deepgram_listen_model(upstream_url: str) -> str: return models[0] if models else DEEPGRAM_LISTEN_DEFAULT_MODEL +def _channel_count(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int): + return None + return value if value >= 1 else None + + +def _results_channel_count(frame: Mapping[str, object]) -> int | None: + channel_index: Final = frame.get("channel_index") + if not isinstance(channel_index, list) or len(channel_index) != 2: + return None + return _channel_count(channel_index[1]) + + +def _declared_channel_count(upstream_url: str) -> int | None: + declared: Final = parse_qs(urlparse(upstream_url).query).get("channels") + if not declared or not declared[0].isdigit(): + return None + return _channel_count(int(declared[0])) + + +def deepgram_listen_channel_count(websocket_messages: Sequence[Mapping[str, object]], upstream_url: str) -> int: + metadata_channels: Final = tuple( + channels + for frame in websocket_messages + if frame.get("type") == "Metadata" + if (channels := _channel_count(frame.get("channels"))) is not None + ) + if metadata_channels: + return metadata_channels[-1] + results_channels: Final = tuple( + channels + for frame in websocket_messages + if frame.get("type") == "Results" + if (channels := _results_channel_count(frame)) is not None + ) + if results_channels: + return max(results_channels) + return _declared_channel_count(upstream_url) or 1 + + def _seconds(value: object) -> float | None: if isinstance(value, bool) or not isinstance(value, (int, float)): return None diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py index a5fea7c8020..0a8d7ff3d33 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.deepgram.common_utils import ( deepgram_listen_audio_seconds, + deepgram_listen_channel_count, deepgram_listen_model, deepgram_listen_transcript, ) @@ -45,8 +46,10 @@ class DeepgramListenPassthroughLoggingHandler: ) -> PassThroughEndpointLoggingTypedDict: model: Final = deepgram_listen_model(upstream_url) audio_seconds: Final = deepgram_listen_audio_seconds(websocket_messages) + channels: Final = deepgram_listen_channel_count(websocket_messages, upstream_url) + billed_seconds: Final = audio_seconds * channels response: Final = TranscriptionResponse(text=deepgram_listen_transcript(websocket_messages)) - response._hidden_params["audio_transcription_duration"] = audio_seconds # pyright: ignore[reportPrivateUsage] # the cost calculator reads the billed duration off the response's hidden params + response._hidden_params["audio_transcription_duration"] = billed_seconds # pyright: ignore[reportPrivateUsage] # the cost calculator reads the billed duration off the response's hidden params response_cost: Final = _audio_cost(response, model) response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads a precomputed cost off the response's hidden params @@ -56,9 +59,10 @@ class DeepgramListenPassthroughLoggingHandler: logging_obj.model_call_details["custom_llm_provider"] = provider # rebind-ok: same shared logging object logging_obj.model_call_details["response_cost"] = response_cost # rebind-ok: same shared logging object verbose_proxy_logger.debug( - "Deepgram listen passthrough cost tracking: model %s, audio seconds %s, cost %s", + "Deepgram listen passthrough cost tracking: model %s, audio seconds %s, channels %s, cost %s", model, audio_seconds, + channels, response_cost, ) logging_result: Final[PassThroughEndpointLoggingTypedDict] = { diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py index 65fbf7c7870..d7a38f2ca5a 100644 --- a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py @@ -8,6 +8,7 @@ import litellm from litellm.llms.deepgram.common_utils import ( deepgram_listen_audio_seconds, deepgram_listen_callback_params, + deepgram_listen_channel_count, deepgram_listen_model, deepgram_listen_transcript, deepgram_listen_websocket_target, @@ -16,18 +17,25 @@ from litellm.llms.deepgram.common_utils import ( NOVA_3_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" -def _results(start: object, duration: object, transcript: str = "", is_final: object = True) -> dict[str, object]: +def _results( + start: object, + duration: object, + transcript: str = "", + is_final: object = True, + channel_index: object = (0, 1), +) -> dict[str, object]: return { "type": "Results", "start": start, "duration": duration, "is_final": is_final, + "channel_index": list(channel_index) if isinstance(channel_index, tuple) else channel_index, "channel": {"alternatives": [{"transcript": transcript, "confidence": 0.9}]}, } -def _metadata(duration: object) -> dict[str, object]: - return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": 1} +def _metadata(duration: object, channels: object = 1) -> dict[str, object]: + return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": channels} @pytest.mark.parametrize( @@ -107,6 +115,62 @@ def test_deepgram_listen_audio_seconds(frames: Sequence[Mapping[str, object]], e assert deepgram_listen_audio_seconds(frames) == expected_seconds +@pytest.mark.parametrize( + ("frames", "upstream_url", "expected_channels"), + [ + pytest.param((_results(0.0, 2.0), _metadata(6.25)), NOVA_3_URL, 1, id="mono"), + pytest.param((_results(0.0, 2.0, channel_index=(0, 2)), _metadata(6.25, 2)), NOVA_3_URL, 2, id="stereo"), + pytest.param((_metadata(1.0, 3), _metadata(1.0, 5)), NOVA_3_URL, 5, id="last metadata wins"), + pytest.param( + (_metadata(1.0, 20), _results(0.0, 1.0, channel_index=(1, 2))), + NOVA_3_URL, + 20, + id="metadata beats channel_index", + ), + pytest.param( + (_results(0.0, 1.0, channel_index=(0, 2)), _results(0.0, 1.0, channel_index=(3, 4))), + NOVA_3_URL, + 4, + id="widest channel_index without metadata", + ), + pytest.param( + (_results(0.0, 1.0, channel_index=(0, 2)),), + f"{NOVA_3_URL}&channels=7&multichannel=true", + 2, + id="frames beat the declared query", + ), + pytest.param((), f"{NOVA_3_URL}&channels=7&multichannel=true", 7, id="declared query when no frames"), + pytest.param((), f"{NOVA_3_URL}&channels=0", 1, id="zero declared channels"), + pytest.param((), f"{NOVA_3_URL}&channels=-2", 1, id="negative declared channels"), + pytest.param((), f"{NOVA_3_URL}&channels=two", 1, id="non numeric declared channels"), + pytest.param((), NOVA_3_URL, 1, id="nothing declared"), + pytest.param((_metadata(1.0, "2"), _metadata(1.0, True), _metadata(1.0, 0)), NOVA_3_URL, 1, id="bad metadata"), + pytest.param((_metadata(1.0, 3), _metadata(1.0, True)), NOVA_3_URL, 3, id="boolean does not shadow a count"), + pytest.param((_metadata(1.0, 2.0), _metadata(1.0, -1)), NOVA_3_URL, 1, id="float and negative metadata"), + pytest.param( + (_metadata(1.0, 2), {**_results(0.0, 1.0), "channels": 9}, {"type": "UtteranceEnd", "channels": 11}), + NOVA_3_URL, + 2, + id="channels on non metadata frames ignored", + ), + pytest.param( + ( + _results(0.0, 1.0, channel_index=[0]), + _results(0.0, 1.0, channel_index=(0, "2")), + _results(0.0, 1.0, channel_index=(0, 0)), + ), + NOVA_3_URL, + 1, + id="bad channel_index", + ), + ], +) +def test_deepgram_listen_channel_count( + frames: Sequence[Mapping[str, object]], upstream_url: str, expected_channels: int +): + assert deepgram_listen_channel_count(frames, upstream_url) == expected_channels + + def test_deepgram_listen_transcript_joins_final_results_only(): frames = ( _results(0.0, 1.0, "hello wor", is_final=False), diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py index 40f520e344d..944b37f95e3 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py @@ -30,8 +30,8 @@ def _results(start: object, duration: object, transcript: str = "", is_final: ob } -def _metadata(duration: object) -> dict[str, object]: - return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": 1} +def _metadata(duration: object, channels: int = 1) -> dict[str, object]: + return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": channels} @pytest.mark.parametrize( @@ -119,6 +119,34 @@ def test_handler_charges_more_for_more_audio_on_the_same_model(): assert short["kwargs"]["response_cost"] > 0 +def test_handler_bills_every_channel_of_a_multichannel_session(): + """Deepgram bills processed audio per channel (deepgram.com/pricing FAQ, 2026-09-17), so a stereo session must be + charged for twice its wall-clock duration or budgets can be bypassed by requesting more channels.""" + mono = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=(_metadata(30.0),), logging_obj=_logging_obj(), upstream_url=NOVA_3_URL + ) + stereo = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=(_metadata(30.0, channels=2),), + logging_obj=_logging_obj(), + upstream_url=f"{NOVA_3_URL}&multichannel=true&channels=2", + ) + + assert stereo["result"]._hidden_params["audio_transcription_duration"] == 60.0 + assert stereo["kwargs"]["response_cost"] == pytest.approx(2 * mono["kwargs"]["response_cost"]) + assert stereo["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-3", 60.0)) + + +def test_handler_bills_the_declared_channels_when_the_stream_dies_before_any_frame_reports_them(): + handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=(_results(0.0, 10.0, "a"),), + logging_obj=_logging_obj(), + upstream_url=f"{NOVA_3_URL}&multichannel=true&channels=3", + ) + + assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 30.0 + assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-3", 30.0)) + + def test_handler_keeps_the_spend_row_but_no_cost_for_an_unpriced_model(): handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( websocket_messages=(_metadata(12.5),), From 0f3556b1bf505864b8860b41e5d917f38b7e5468 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 19:17:49 +0000 Subject: [PATCH 035/102] refactor(proxy): drop explanatory docstrings from the stream attribution helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 3 +-- litellm/proxy/utils.py | 11 ----------- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9a038ca79ba..849211a2dae 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2576,8 +2576,7 @@ class ProxyBaseLLMRequestProcessing: ) async def refresh_stream_headers() -> Mapping[str, str]: - """`custom_headers` rebuilt once the first chunk is buffered, from `self.data` as the - guardrails left it and for whichever deployment served the stream.""" + """`custom_headers` rebuilt for whichever deployment served the stream.""" return self._stream_response_headers( hidden_params=( get_hidden_params_dict(response) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 21e6ed1d24f..4182a864e4c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -468,9 +468,6 @@ def _record_raising_guardrail(request_data: Mapping[str, object], callback: obje class _UpstreamStreamBoundary(Generic[_T]): - """Remembers the exception the upstream iterator raised, so the wrapper around a - streaming hook can tell a pass-through failure from one the hook raised itself.""" - __slots__ = ("_upstream", "failure") def __init__(self, upstream: AsyncIterable[_T]) -> None: @@ -2548,14 +2545,6 @@ class ProxyLogging: hook: _StreamIteratorHook[_T], request_data: Mapping[str, object], ) -> AsyncGenerator[_T, None]: - """ - Run `hook` over `response` and yield its chunks. If the hook itself raises, - enrich an HTTPException's dict detail with the callback's `guardrail_name` - and `guardrail_mode` and record the callback in `applied_guardrails` before - re-raising. Failures raised by `response` (the provider stream or an inner - layer of the async_post_call_streaming_iterator_hook chain) pass through - untouched, so only the layer that actually raised is attributed. - """ upstream: Final = _UpstreamStreamBoundary(response) try: async for chunk in hook(response=upstream): From 11ec157d71b1d40100ff9d2be778477e0a145238 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:36:03 +0000 Subject: [PATCH 036/102] fix(deepgram): authorize the effective model and price /listen sessions at streaming rates Key auth on the Deepgram WebSocket route now sees the same model the upstream target will carry, so a key restricted to other models can no longer reach nova-3 by leaving model out of the query. user_api_key_auth_websocket keeps its signature and delegates to user_api_key_auth_websocket_for_model, which the Deepgram route calls with deepgram_listen_requested_model Sessions are priced from new deepgram/streaming/* registry rows (nova-3, nova-3-multilingual for language=multi) plus per-minute add-on rows for redact, keyterm, detect_entities and diarize, all read from Deepgram's pricing page on 2026-09-17. Models without a streaming row fall back to their pre-recorded row as before Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/deepgram/common_utils.py | 49 ++++++++++ ...odel_prices_and_context_window_backup.json | 90 ++++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 10 +- .../llm_passthrough_endpoints.py | 10 +- ...gram_listen_passthrough_logging_handler.py | 32 ++++++- model_prices_and_context_window.json | 90 ++++++++++++++++++ .../deepgram/test_deepgram_common_utils.py | 68 ++++++++++++++ ...gram_listen_passthrough_logging_handler.py | 92 +++++++++++++++++-- .../test_deepgram_ws_passthrough_routes.py | 62 ++++++++++++- 9 files changed, 481 insertions(+), 22 deletions(-) diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index 6beb70b4499..4f18918fa71 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -11,12 +11,29 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException _WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"}) DEEPGRAM_LISTEN_CALLBACK_PARAMS: Final = frozenset({"callback", "callback_method"}) +DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX: Final = "streaming/" +DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE: Final = "multi" +DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX: Final = "-multilingual" +DEEPGRAM_LISTEN_ADDON_PRICING_PARAMS: Final = MappingProxyType( + { + "redact": "redact", + "keyterm": "keyterm", + "detect_entities": "detect_entities", + "diarize": "diarize", + "diarize_model": "diarize", + } +) +_DISABLED_PARAM_VALUES: Final = frozenset({"", "false"}) class DeepgramException(BaseLLMException): pass +def deepgram_listen_requested_model(query_string: str) -> str: + return httpx.QueryParams(query_string).get("model") or DEEPGRAM_LISTEN_DEFAULT_MODEL + + def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str: listen_url: Final = httpx.URL(f"{(api_base or DEEPGRAM_DEFAULT_API_BASE).rstrip('/')}/listen") websocket_url: Final = listen_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(listen_url.scheme, listen_url.scheme)) @@ -36,6 +53,38 @@ def deepgram_listen_model(upstream_url: str) -> str: return models[0] if models else DEEPGRAM_LISTEN_DEFAULT_MODEL +def _param_enabled(values: Sequence[str]) -> bool: + return any(value.strip().lower() not in _DISABLED_PARAM_VALUES for value in values) + + +def deepgram_listen_base_pricing_models(upstream_url: str) -> tuple[str, ...]: + """Registry keys to try, in order, for the per-second base rate of a streaming session: the streaming entry for + the language mode Deepgram bills (multilingual when ``language=multi``), then the plain streaming entry, then + the pre-recorded entry for models that have no streaming price of their own.""" + model: Final = deepgram_listen_model(upstream_url) + params: Final = parse_qs(urlparse(upstream_url).query) + streaming: Final = f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{model}" + multilingual: Final = params.get("language", ("",))[-1].strip().lower() == DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE + return ( + (f"{streaming}{DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX}", streaming, model) + if multilingual + else (streaming, model) + ) + + +def deepgram_listen_addon_pricing_models(upstream_url: str) -> tuple[str, ...]: + params: Final = parse_qs(urlparse(upstream_url).query) + return tuple( + sorted( + frozenset( + f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{addon}" + for param, addon in DEEPGRAM_LISTEN_ADDON_PRICING_PARAMS.items() + if _param_enabled(params.get(param, ())) + ) + ) + ) + + def _channel_count(value: object) -> int | None: if isinstance(value, bool) or not isinstance(value, int): return None diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e87a3fec99b..7b7b3c7199c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -20736,6 +20736,96 @@ "/v1/audio/transcriptions" ] }, + "deepgram/streaming/nova-3": { + "input_cost_per_second": 8e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0048/60 seconds = $0.00008000 per second", + "note": "Nova-3 monolingual streaming, pay as you go", + "original_pricing_per_minute": 0.0048 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/nova-3-multilingual": { + "input_cost_per_second": 9.667e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0058/60 seconds = $0.00009667 per second", + "note": "Nova-3 multilingual (language=multi) streaming, pay as you go", + "original_pricing_per_minute": 0.0058 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/redact": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0020/60 seconds = $0.00003333 per second", + "note": "Redaction add-on (redact query param), streaming, pay as you go", + "original_pricing_per_minute": 0.002 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/keyterm": { + "input_cost_per_second": 2.167e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0013/60 seconds = $0.00002167 per second", + "note": "Keyterm Prompting add-on (keyterm query param), streaming, pay as you go", + "original_pricing_per_minute": 0.0013 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/detect_entities": { + "input_cost_per_second": 2.833e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0017/60 seconds = $0.00002833 per second", + "note": "Entity Detection add-on (detect_entities query param), streaming, pay as you go", + "original_pricing_per_minute": 0.0017 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/diarize": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0020/60 seconds = $0.00003333 per second", + "note": "Speaker Diarization add-on (diarize / diarize_model query params), streaming, pay as you go", + "original_pricing_per_minute": 0.002 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, "deepgram/whisper": { "input_cost_per_second": 0.0001, "litellm_provider": "deepgram", diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4cbd4213463..6842d500a82 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -629,9 +629,11 @@ def _apply_budget_limits_to_end_user_params( verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) -async def user_api_key_auth_websocket(websocket: WebSocket): - # Accept the WebSocket connection +async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth: + return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model")) + +async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str | None) -> UserAPIKeyAuth: ws_scope: Final = websocket.scope or {} scope_headers: Final = list(ws_scope.get("headers") or []) # ``get_request_route`` falls back to ``request.url.path`` when @@ -651,10 +653,6 @@ async def user_api_key_auth_websocket(websocket: WebSocket): request._url = websocket.url - query_params: Final = websocket.query_params - - model: Final = query_params.get("model") - async def return_body(): return _realtime_request_body(model) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7f9e0169fb2..0bd3eb941b9 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -38,6 +38,7 @@ from litellm.llms.azure.passthrough.transformation import foreign_azure_deployme from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.deepgram.common_utils import ( deepgram_listen_callback_params, + deepgram_listen_requested_model, deepgram_listen_websocket_target, ) from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path @@ -52,6 +53,7 @@ from litellm.proxy.auth.user_api_key_auth import ( is_no_auth_dev_mode, user_api_key_auth, user_api_key_auth_websocket, + user_api_key_auth_websocket_for_model, ) from litellm.proxy.common_request_processing import open_sse_before_first_byte from litellm.proxy.common_utils.http_parsing_utils import ( @@ -2700,11 +2702,17 @@ _DEEPGRAM_WS_MISSING_KEY_REASON: Final = ( _DEEPGRAM_WS_CALLBACK_REASON: Final = "Deepgram callback delivery is not supported through the proxy: remove {params}" +async def deepgram_listen_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth: + return await user_api_key_auth_websocket_for_model( + websocket, model=deepgram_listen_requested_model(websocket.url.query) + ) + + @router.websocket("/deepgram/v1/listen") @router.websocket("/deepgram/listen") async def deepgram_listen_websocket_route( websocket: WebSocket, - user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(deepgram_listen_user_api_key_auth)], relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], ) -> None: deepgram_api_key: Final = passthrough_endpoint_router.get_credentials( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py index 0a8d7ff3d33..fc0af400477 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py @@ -7,7 +7,9 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.deepgram.common_utils import ( + deepgram_listen_addon_pricing_models, deepgram_listen_audio_seconds, + deepgram_listen_base_pricing_models, deepgram_listen_channel_count, deepgram_listen_model, deepgram_listen_transcript, @@ -18,19 +20,39 @@ from litellm.types.utils import TranscriptionResponse DEEPGRAM_LISTEN_ROUTE_SUFFIX: Final = "/listen" -def _audio_cost(response: TranscriptionResponse, model: str) -> float | None: +def _registry_cost(response: TranscriptionResponse, pricing_model: str) -> float | None: try: return litellm.completion_cost( completion_response=response, - model=model, + model=pricing_model, custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value, call_type="transcription", ) - except Exception as e: # noqa: BLE001 # an unpriced model must not lose the spend row, only its cost - verbose_proxy_logger.warning("Deepgram listen passthrough: no pricing for model '%s': %s", model, e) + except Exception as e: # noqa: BLE001 # an unpriced entry must not lose the spend row, only its cost + verbose_proxy_logger.debug("Deepgram listen passthrough: no registry price for '%s': %s", pricing_model, e) return None +def _audio_cost(response: TranscriptionResponse, upstream_url: str) -> float | None: + base_cost: Final = next( + ( + cost + for pricing_model in deepgram_listen_base_pricing_models(upstream_url) + if (cost := _registry_cost(response, pricing_model)) is not None + ), + None, + ) + if base_cost is None: + verbose_proxy_logger.warning( + "Deepgram listen passthrough: no pricing for model '%s'", deepgram_listen_model(upstream_url) + ) + return None + addon_costs: Final = tuple( + _registry_cost(response, pricing_model) for pricing_model in deepgram_listen_addon_pricing_models(upstream_url) + ) + return base_cost + sum(cost for cost in addon_costs if cost is not None) + + class DeepgramListenPassthroughLoggingHandler: @staticmethod def is_deepgram_listen_route(url_route: str) -> bool: @@ -50,7 +72,7 @@ class DeepgramListenPassthroughLoggingHandler: billed_seconds: Final = audio_seconds * channels response: Final = TranscriptionResponse(text=deepgram_listen_transcript(websocket_messages)) response._hidden_params["audio_transcription_duration"] = billed_seconds # pyright: ignore[reportPrivateUsage] # the cost calculator reads the billed duration off the response's hidden params - response_cost: Final = _audio_cost(response, model) + response_cost: Final = _audio_cost(response, upstream_url) response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads a precomputed cost off the response's hidden params provider: Final = litellm.LlmProviders.DEEPGRAM.value diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e87a3fec99b..7b7b3c7199c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -20736,6 +20736,96 @@ "/v1/audio/transcriptions" ] }, + "deepgram/streaming/nova-3": { + "input_cost_per_second": 8e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0048/60 seconds = $0.00008000 per second", + "note": "Nova-3 monolingual streaming, pay as you go", + "original_pricing_per_minute": 0.0048 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/nova-3-multilingual": { + "input_cost_per_second": 9.667e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0058/60 seconds = $0.00009667 per second", + "note": "Nova-3 multilingual (language=multi) streaming, pay as you go", + "original_pricing_per_minute": 0.0058 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/redact": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0020/60 seconds = $0.00003333 per second", + "note": "Redaction add-on (redact query param), streaming, pay as you go", + "original_pricing_per_minute": 0.002 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/keyterm": { + "input_cost_per_second": 2.167e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0013/60 seconds = $0.00002167 per second", + "note": "Keyterm Prompting add-on (keyterm query param), streaming, pay as you go", + "original_pricing_per_minute": 0.0013 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/detect_entities": { + "input_cost_per_second": 2.833e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0017/60 seconds = $0.00002833 per second", + "note": "Entity Detection add-on (detect_entities query param), streaming, pay as you go", + "original_pricing_per_minute": 0.0017 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/diarize": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0020/60 seconds = $0.00003333 per second", + "note": "Speaker Diarization add-on (diarize / diarize_model query params), streaming, pay as you go", + "original_pricing_per_minute": 0.002 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://deepgram.com/pricing", + "supported_endpoints": [ + "/v1/listen" + ] + }, "deepgram/whisper": { "input_cost_per_second": 0.0001, "litellm_provider": "deepgram", diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py index d7a38f2ca5a..802a3795ffe 100644 --- a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py @@ -6,10 +6,13 @@ import pytest import litellm from litellm.llms.deepgram.common_utils import ( + deepgram_listen_addon_pricing_models, deepgram_listen_audio_seconds, + deepgram_listen_base_pricing_models, deepgram_listen_callback_params, deepgram_listen_channel_count, deepgram_listen_model, + deepgram_listen_requested_model, deepgram_listen_transcript, deepgram_listen_websocket_target, ) @@ -195,3 +198,68 @@ def test_deepgram_listen_transcript_joins_final_results_only(): ) def test_deepgram_listen_model_comes_from_the_upstream_query(upstream_url: str, expected_model: str): assert deepgram_listen_model(upstream_url) == expected_model + + +@pytest.mark.parametrize( + "query_string", + ["model=nova-2&language=en", "language=en", "model=&language=en", "", "model=nova-3-medical"], +) +def test_requested_model_is_the_model_the_upstream_target_will_carry(query_string: str): + """Authorization runs against ``deepgram_listen_requested_model``; the upstream URL is built separately, so the + two must always agree or a key could be authorized for one model and reach another.""" + target: Final = deepgram_listen_websocket_target(None, query_string) + assert deepgram_listen_requested_model(query_string) == deepgram_listen_model(target) + + +@pytest.mark.parametrize( + ("upstream_url", "expected"), + [ + pytest.param(NOVA_3_URL, ("streaming/nova-3", "nova-3"), id="monolingual"), + pytest.param(f"{NOVA_3_URL}&language=en", ("streaming/nova-3", "nova-3"), id="explicit language"), + pytest.param( + f"{NOVA_3_URL}&language=multi", + ("streaming/nova-3-multilingual", "streaming/nova-3", "nova-3"), + id="multilingual", + ), + pytest.param( + f"{NOVA_3_URL}&language=MULTI", + ("streaming/nova-3-multilingual", "streaming/nova-3", "nova-3"), + id="multilingual any case", + ), + pytest.param( + "wss://api.deepgram.com/v1/listen?model=nova-2&language=multi", + ("streaming/nova-2-multilingual", "streaming/nova-2", "nova-2"), + id="other model", + ), + ], +) +def test_deepgram_listen_base_pricing_models(upstream_url: str, expected: tuple[str, ...]): + assert deepgram_listen_base_pricing_models(upstream_url) == expected + + +@pytest.mark.parametrize( + ("upstream_url", "expected"), + [ + pytest.param(NOVA_3_URL, (), id="no add-ons"), + pytest.param(f"{NOVA_3_URL}&redact=pci", ("streaming/redact",), id="redact"), + pytest.param(f"{NOVA_3_URL}&redact=pci&redact=ssn", ("streaming/redact",), id="repeated redact once"), + pytest.param(f"{NOVA_3_URL}&keyterm=a&keyterm=b", ("streaming/keyterm",), id="keyterm"), + pytest.param(f"{NOVA_3_URL}&detect_entities=true", ("streaming/detect_entities",), id="detect_entities"), + pytest.param(f"{NOVA_3_URL}&diarize=true", ("streaming/diarize",), id="diarize"), + pytest.param(f"{NOVA_3_URL}&diarize_model=v1", ("streaming/diarize",), id="diarize_model"), + pytest.param(f"{NOVA_3_URL}&diarize=true&diarize_model=latest", ("streaming/diarize",), id="diarize both once"), + pytest.param(f"{NOVA_3_URL}&detect_entities=false&diarize=FALSE&redact=", (), id="disabled"), + pytest.param( + f"{NOVA_3_URL}&detect_entities=false&detect_entities=true", + ("streaming/detect_entities",), + id="any enabling value wins", + ), + pytest.param( + f"{NOVA_3_URL}&diarize=true&redact=pci&keyterm=x&detect_entities=true", + ("streaming/detect_entities", "streaming/diarize", "streaming/keyterm", "streaming/redact"), + id="all, sorted", + ), + ], +) +def test_deepgram_listen_addon_pricing_models(upstream_url: str, expected: tuple[str, ...]): + assert deepgram_listen_addon_pricing_models(upstream_url) == expected diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py index 944b37f95e3..a6133a7cf99 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py @@ -19,6 +19,8 @@ from litellm.types.utils import StandardLoggingPayload, TranscriptionResponse NOVA_3_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" +pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map") + def _results(start: object, duration: object, transcript: str = "", is_final: object = True) -> dict[str, object]: return { @@ -63,13 +65,22 @@ def _logging_obj(call_id: str = "call-dg") -> LiteLLMLoggingObj: ) -def _registry_cost(model: str, seconds: float) -> float: +def _registry_cost(pricing_model: str, seconds: float) -> float: """Derives the expected charge from the live cost map rather than pinning a vendor price.""" - per_second: Final = litellm.model_cost[f"deepgram/{model}"]["input_cost_per_second"] + per_second: Final = litellm.model_cost[f"deepgram/{pricing_model}"]["input_cost_per_second"] assert per_second > 0 return per_second * seconds +def _cost(upstream_url: str, *frames: dict[str, object]) -> float: + handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=frames, logging_obj=_logging_obj(), upstream_url=upstream_url + ) + response_cost = handler_result["kwargs"]["response_cost"] + assert isinstance(response_cost, float) + return response_cost + + def test_handler_bills_metadata_duration_at_the_registry_rate_and_names_the_model(): frames = (_results(0.0, 5.0, "first sentence"), _results(5.0, 7.5, "second sentence"), _metadata(12.5)) logging_obj = _logging_obj() @@ -85,15 +96,78 @@ def test_handler_bills_metadata_duration_at_the_registry_rate_and_names_the_mode assert isinstance(result, TranscriptionResponse) assert result.text == "first sentence second sentence" assert result._hidden_params["audio_transcription_duration"] == 12.5 - assert result._hidden_params["response_cost"] == pytest.approx(_registry_cost("nova-3", 12.5)) - assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-3", 12.5)) + assert result._hidden_params["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 12.5)) + assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 12.5)) assert handler_result["kwargs"]["model"] == "nova-3" assert handler_result["kwargs"]["custom_llm_provider"] == "deepgram" assert handler_result["kwargs"]["litellm_params"] == {"metadata": {}} assert logging_obj.model == "nova-3" assert logging_obj.model_call_details["model"] == "nova-3" assert logging_obj.model_call_details["custom_llm_provider"] == "deepgram" - assert logging_obj.model_call_details["response_cost"] == pytest.approx(_registry_cost("nova-3", 12.5)) + assert logging_obj.model_call_details["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 12.5)) + + +def test_handler_bills_streaming_not_prerecorded_rates(): + """Deepgram prices /v1/listen over a WebSocket separately from pre-recorded transcription, so the streaming entry + must be the one charged; the two registry rows only need to differ for this to matter, whatever their values.""" + streaming = litellm.model_cost["deepgram/streaming/nova-3"]["input_cost_per_second"] + prerecorded = litellm.model_cost["deepgram/nova-3"]["input_cost_per_second"] + assert streaming != prerecorded + + assert _cost(NOVA_3_URL, _metadata(60.0)) == pytest.approx(60.0 * streaming) + + +def test_handler_bills_multilingual_streaming_when_language_is_multi(): + monolingual = _cost(NOVA_3_URL, _metadata(60.0)) + multilingual = _cost(f"{NOVA_3_URL}&language=multi", _metadata(60.0)) + + assert multilingual == pytest.approx(_registry_cost("streaming/nova-3-multilingual", 60.0)) + assert multilingual > monolingual + + +@pytest.mark.parametrize( + ("query", "addons"), + [ + pytest.param("redact=pci", ("redact",), id="redaction"), + pytest.param("redact=pci&redact=numbers", ("redact",), id="redaction counted once"), + pytest.param("keyterm=LiteLLM&keyterm=Deepgram", ("keyterm",), id="keyterm prompting"), + pytest.param("detect_entities=true", ("detect_entities",), id="entity detection"), + pytest.param("diarize=true", ("diarize",), id="diarization"), + pytest.param("diarize_model=v1", ("diarize",), id="diarization via diarize_model"), + pytest.param("diarize=true&diarize_model=v1", ("diarize",), id="diarization counted once"), + pytest.param( + "redact=pci&keyterm=x&detect_entities=true&diarize=true", + ("redact", "keyterm", "detect_entities", "diarize"), + id="every add-on", + ), + pytest.param("detect_entities=false&diarize=False&redact=", (), id="disabled add-ons cost nothing"), + ], +) +def test_handler_adds_each_priced_add_on_once_on_top_of_the_base_rate(query: str, addons: tuple[str, ...]): + base = _cost(NOVA_3_URL, _metadata(60.0)) + expected = base + sum(_registry_cost(f"streaming/{addon}", 60.0) for addon in addons) + + assert _cost(f"{NOVA_3_URL}&{query}", _metadata(60.0)) == pytest.approx(expected) + + +def test_handler_add_ons_scale_with_channels_like_the_base_rate(): + stereo_plain = _cost(f"{NOVA_3_URL}&channels=2", _metadata(60.0, channels=2)) + stereo_redacted = _cost(f"{NOVA_3_URL}&channels=2&redact=pci", _metadata(60.0, channels=2)) + + assert stereo_redacted - stereo_plain == pytest.approx(_registry_cost("streaming/redact", 120.0)) + + +def test_handler_falls_back_to_the_prerecorded_rate_for_a_model_without_a_streaming_entry(): + assert "deepgram/streaming/nova-2" not in litellm.model_cost + + handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=(_metadata(60.0),), + logging_obj=_logging_obj(), + upstream_url="wss://api.deepgram.com/v1/listen?model=nova-2", + ) + + assert handler_result["kwargs"]["model"] == "nova-2" + assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-2", 60.0)) def test_handler_falls_back_to_results_frames_when_the_stream_ends_without_metadata(): @@ -104,7 +178,7 @@ def test_handler_falls_back_to_results_frames_when_the_stream_ends_without_metad ) assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 72.5 - assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-3", 72.5)) + assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 72.5)) def test_handler_charges_more_for_more_audio_on_the_same_model(): @@ -133,7 +207,7 @@ def test_handler_bills_every_channel_of_a_multichannel_session(): assert stereo["result"]._hidden_params["audio_transcription_duration"] == 60.0 assert stereo["kwargs"]["response_cost"] == pytest.approx(2 * mono["kwargs"]["response_cost"]) - assert stereo["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-3", 60.0)) + assert stereo["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 60.0)) def test_handler_bills_the_declared_channels_when_the_stream_dies_before_any_frame_reports_them(): @@ -144,7 +218,7 @@ def test_handler_bills_the_declared_channels_when_the_stream_dies_before_any_fra ) assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 30.0 - assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-3", 30.0)) + assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 30.0)) def test_handler_keeps_the_spend_row_but_no_cost_for_an_unpriced_model(): @@ -226,6 +300,6 @@ async def test_success_handler_dispatches_deepgram_listen_and_logs_duration_base payload = capturing_logger.payloads[0] assert payload["model"] == "nova-3" assert payload["custom_llm_provider"] == "deepgram" - assert payload["response_cost"] == pytest.approx(_registry_cost("nova-3", 20.0)) + assert payload["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 20.0)) assert payload["metadata"]["user_api_key_team_id"] == "team-stt" assert payload["id"] == "call-dg-e2e" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 5ea2b0b8ab9..1c21b75804e 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -1,10 +1,11 @@ """Deepgram ``/v1/listen`` passthrough WebSocket route: registration, auth, credential injection, target URL.""" +import asyncio from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI @@ -12,13 +13,16 @@ from fastapi.testclient import TestClient from starlette.routing import WebSocketRoute from starlette.websockets import WebSocketDisconnect +from litellm.caching.dual_cache import DualCache from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import _cache_key_object from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _websocket_relay, deepgram_listen_websocket_route, router, ) +from litellm.proxy.utils import hash_token GET_CREDENTIALS: Final = ( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials" @@ -302,6 +306,62 @@ def test_deepgram_listen_authenticates_the_litellm_key_and_relays_to_deepgram(mo ] +async def _cache_restricted_key(virtual_key: str, models: list[str]) -> DualCache: + cache = DualCache() + await _cache_key_object( + hashed_token=hash_token(virtual_key), + user_api_key_obj=UserAPIKeyAuth(token=hash_token(virtual_key), models=models), + user_api_key_cache=cache, + proxy_logging_obj=None, + ) + return cache + + +@pytest.mark.parametrize( + ("query", "expect_relay"), + [ + pytest.param("model=nova-2", True, id="allowed model named"), + pytest.param("model=nova-3", False, id="denied model named"), + pytest.param("", False, id="model omitted, default denied"), + pytest.param("model=&language=en", False, id="model blank, default denied"), + ], +) +def test_deepgram_listen_authorizes_the_model_it_will_actually_send_upstream(query, expect_relay, monkeypatch): + """A key allowed only ``nova-2`` must not reach ``nova-3`` by leaving ``model`` out and letting the proxy fill + in its default: the real key auth path must see the same model the upstream target will carry.""" + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + cache = asyncio.run(_cache_restricted_key("sk-only-nova-2", ["nova-2"])) + relay = _FakeRelay() + client = TestClient(_app_with_relay(relay)) + + with ( + patch(GET_CREDENTIALS, return_value="dg-provider-key"), + patch.multiple( # test-quality-ok: the real key auth path reads these proxy_server globals and has no injection seam + "litellm.proxy.proxy_server", + master_key="sk-master", + prisma_client=MagicMock(), + user_api_key_cache=cache, + llm_model_list=None, + llm_router=None, + ), + ): + if expect_relay: + with client.websocket_connect( + f"/deepgram/v1/listen?{query}", headers={"Authorization": "Bearer sk-only-nova-2"} + ): + pass + assert [call.target for call in relay.calls] == [f"wss://api.deepgram.com/v1/listen?{query}"] + return + with pytest.raises(WebSocketDisconnect) as disconnect: + with client.websocket_connect( + f"/deepgram/v1/listen?{query}", headers={"Authorization": "Bearer sk-only-nova-2"} + ): + pass + + assert disconnect.value.code == 1008 + assert relay.calls == [] + + def test_deepgram_listen_echoes_the_browser_subprotocol_that_carries_the_litellm_key(monkeypatch): """Browsers cannot set headers, so they send the key as a subprotocol and abort the handshake unless the server echoes that subprotocol back; the key itself must still stay off the upstream connection.""" From 3084d2af31794f552d1584558fc1efadfdcc7dc9 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:51:50 +0000 Subject: [PATCH 037/102] test(utils): allow /v1/listen in the registry supported_endpoints schema The deepgram/streaming/* rows added for the Deepgram WebSocket passthrough declare /v1/listen as their endpoint, so the registry validation test needs it in the enum, the same way /vertex_ai/live was added for that passthrough Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index f219f26b353..77336836e04 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -945,6 +945,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "/v1/audio/speech", "/v1/ocr", "/vertex_ai/live", + "/v1/listen", "/v1beta/interactions", ], }, From 6e7c3f68a1f3cebe75c2e019490db5cfb143f56f Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 21:06:07 +0000 Subject: [PATCH 038/102] test(deepgram): pin litellm.max_budget to zero in the model authorization route test Under xdist the per-test litellm reload is skipped, so a leaked max_budget from another proxy test sent the real key auth path into the global spend lookup, which the MagicMock prisma client cannot await Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_deepgram_ws_passthrough_routes.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 1c21b75804e..85baaacf1e9 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -13,6 +13,7 @@ from fastapi.testclient import TestClient from starlette.routing import WebSocketRoute from starlette.websockets import WebSocketDisconnect +import litellm from litellm.caching.dual_cache import DualCache from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth @@ -330,6 +331,7 @@ def test_deepgram_listen_authorizes_the_model_it_will_actually_send_upstream(que """A key allowed only ``nova-2`` must not reach ``nova-3`` by leaving ``model`` out and letting the proxy fill in its default: the real key auth path must see the same model the upstream target will carry.""" monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setattr(litellm, "max_budget", 0.0) cache = asyncio.run(_cache_restricted_key("sk-only-nova-2", ["nova-2"])) relay = _FakeRelay() client = TestClient(_app_with_relay(relay)) From 9f6e5242853a601670774871073d9175d15edfd0 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 01:43:03 +0000 Subject: [PATCH 039/102] fix(deepgram): ignore zero-duration Metadata frames when billing streamed audio Deepgram sends a Metadata frame with duration 0 on connect. When the closing Metadata frame is not collected before the socket closes, that handshake frame used to become the billed duration and the session logged zero spend. Only a positive Metadata duration is treated as authoritative now; otherwise the furthest Results end time is billed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/deepgram/common_utils.py | 2 +- .../llms/deepgram/test_deepgram_common_utils.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index 4f18918fa71..ede5b4f157d 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -152,7 +152,7 @@ def deepgram_listen_audio_seconds(websocket_messages: Sequence[Mapping[str, obje duration for frame in websocket_messages if frame.get("type") == "Metadata" - if (duration := _seconds(frame.get("duration"))) is not None + if (duration := _seconds(frame.get("duration"))) is not None and duration > 0 ) if metadata_durations: return metadata_durations[-1] diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py index 802a3795ffe..8b61d5bffa0 100644 --- a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py @@ -105,6 +105,12 @@ def test_deepgram_listen_callback_params(query_string: str, expected: tuple[str, pytest.param((_metadata(4.0), _results(0.0, 9.0), _metadata(5.5)), 5.5, id="last metadata wins"), pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _results(1.0, 1.0)), 5.5, id="furthest results end"), pytest.param((_results(0.0, 0.0), _metadata(0.0)), 0.0, id="zero metadata is a real zero"), + pytest.param( + (_metadata(0.0), _results(0.0, 2.0), _results(2.0, 3.5)), + 5.5, + id="handshake metadata zero does not hide streamed results", + ), + pytest.param((_metadata(0.0), _results(0.0, 2.0), _metadata(0.0)), 2.0, id="only zero metadata frames"), pytest.param((_results(0.0, 1.5), _metadata("6.25")), 1.5, id="string metadata is ignored"), pytest.param((_results(0.0, 1.5), _metadata(True)), 1.5, id="boolean metadata is ignored"), pytest.param((_results(0.0, 1.5), _metadata(-3.0)), 1.5, id="negative metadata is ignored"), From a187a5bfc6e64de75bd59aca3103bf2059deeb45 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 07:54:01 +0000 Subject: [PATCH 040/102] fix(proxy): keep requested model guardrails and key disable_fallbacks on rate-limit fallback The local rate-limit fallback path introduced in #40596 restores the request from a snapshot taken before the pre-call pass, so the guardrails resolved for the requested model were dropped when another deployment was selected, and the raw request-body disable_fallbacks field gated the fallback before the key-level disable_fallbacks override had been applied Carry the requested model's merged guardrail list onto every fallback pass and read the effective disable_fallbacks value after the pre-call pass Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 23 +++++- .../proxy/test_common_request_processing.py | 70 ++++++++++++++++++- 2 files changed, 87 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f650b6d0b28..607ce3441e9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -94,7 +94,11 @@ from litellm.proxy.common_utils.sse_keepalive import ( from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.route_llm_request import route_request -from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails +from litellm.proxy.utils import ( + ProxyLogging, + _check_and_merge_model_level_guardrails, + _merge_guardrails_with_existing, +) from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_utils.common_utils import resolve_model_group_alias @@ -1934,6 +1938,7 @@ class ProxyBaseLLMRequestProcessing: user_api_base: str | None = None, model: str | None = None, llm_router: Router | None = None, + requested_model_guardrails: list | None = None, ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks @@ -2098,7 +2103,11 @@ class ProxyBaseLLMRequestProcessing: # could otherwise spoof an unguarded model_info.id while requesting # a guarded alias and bypass guardrails (veria-ai HIGH on #29654). self.data = _check_and_merge_model_level_guardrails( - data=self.data, + data=( + self.data + if requested_model_guardrails is None + else _merge_guardrails_with_existing(data=self.data, model_level_guardrails=requested_model_guardrails) + ), llm_router=llm_router, trust_client_model_info=False, ) @@ -2163,7 +2172,7 @@ class ProxyBaseLLMRequestProcessing: configured_fallbacks: Final = ( self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict) - if llm_router is not None and not self.data.get("disable_fallbacks") + if llm_router is not None else None ) pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None @@ -2208,6 +2217,7 @@ class ProxyBaseLLMRequestProcessing: original_model, fallback_models, ) + requested_model_guardrails: Final = self._request_guardrails(rate_limited_data) try: for fallback_model in fallback_models: @@ -2231,6 +2241,7 @@ class ProxyBaseLLMRequestProcessing: model=fallback_model, route_type=route_type, llm_router=llm_router, + requested_model_guardrails=requested_model_guardrails, ) except ProxyRateLimitError: continue @@ -2241,6 +2252,12 @@ class ProxyBaseLLMRequestProcessing: self.data = rate_limited_data raise original_exc + @staticmethod + def _request_guardrails(data: dict) -> list | None: + metadata: Final = data.get("metadata") + guardrails: Final = metadata.get("guardrails") if isinstance(metadata, dict) else None + return guardrails if isinstance(guardrails, list) else None + @staticmethod def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None: key_router_settings: Final = user_api_key_dict.router_settings diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index d465deace15..f440711ac29 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6711,6 +6711,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: monkeypatch: pytest.MonkeyPatch, user_api_key_dict: ProxyUserAPIKeyAuth, fallbacks: list[dict[str, list[str]]], + model_guardrails: dict[str, list[str]] | None = None, ) -> tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]]: """Real v3 limiter (the default ``parallel_request_limiter``) wired in through the ``proxy_logging_obj`` seam, so ``common_processing_pre_call_logic`` runs for real: @@ -6738,9 +6739,17 @@ class TestPreCallWithFallbacksOnLocalRateLimit: proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=run_limiter) + guardrails_by_group = model_guardrails or {} router = litellm.Router( model_list=[ - {"model_name": group, "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}} + { + "model_name": group, + "litellm_params": { + "model": "openai/gpt-4.1-nano", + "api_key": "fake", + **({"guardrails": guardrails_by_group[group]} if group in guardrails_by_group else {}), + }, + } for chain in fallbacks for group in (*chain.keys(), *(m for models in chain.values() for m in models)) ], @@ -6752,7 +6761,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: def _otel_key( rpm_limit: int | None = None, model_rpm_limit: dict[str, int] | None = None, - disable_fallbacks: bool = False, + disable_fallbacks: bool | None = None, ) -> ProxyUserAPIKeyAuth: from opentelemetry.sdk.trace import TracerProvider @@ -6763,7 +6772,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: rpm_limit=rpm_limit, metadata={ **({"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}), - **({"disable_fallbacks": True} if disable_fallbacks else {}), + **({"disable_fallbacks": disable_fallbacks} if disable_fallbacks is not None else {}), }, ) @@ -6906,6 +6915,61 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert exc_info.value.status_code == 429 assert rig[3] == [primary_model, primary_model] + @pytest.mark.asyncio + async def test_key_metadata_disable_fallbacks_false_overrides_request_body(self, monkeypatch: pytest.MonkeyPatch): + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + key = self._otel_key(model_rpm_limit={primary_model: 1}, disable_fallbacks=False) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request = { + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "disable_fallbacks": True, + } + + await self._pre_call(dict(request), key, rig) + _, (data, _) = await self._pre_call(dict(request), key, rig) + + assert data["model"] == fallback_model + assert data["disable_fallbacks"] is False + assert rig[3] == [primary_model, primary_model, fallback_model] + + @pytest.mark.asyncio + async def test_fallback_keeps_requested_model_guardrails(self, monkeypatch: pytest.MonkeyPatch): + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + guardrail = "pii-guard-for-primary" + key = self._otel_key(model_rpm_limit={primary_model: 1}) + rig = self._v3_limiter_rig( + monkeypatch, key, [{primary_model: [fallback_model]}], model_guardrails={primary_model: [guardrail]} + ) + run_limiter = rig[0].pre_call_hook + + async def limiter_then_guardrail( + user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str + ) -> dict[str, object]: + limited = await run_limiter(user_api_key_dict=user_api_key_dict, data=data, call_type=call_type) + if guardrail not in (limited["metadata"].get("guardrails") or []): + return limited + return { + **limited, + "messages": [ + {**m, "content": str(m["content"]).replace("123-45-6789", "[REDACTED-SSN]")} + for m in limited["messages"] + ], + } + + rig[0].pre_call_hook = AsyncMock(side_effect=limiter_then_guardrail) + request = {"model": primary_model, "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]} + + await self._pre_call(dict(request), key, rig) + _, (data, _) = await self._pre_call(dict(request), key, rig) + + assert data["model"] == fallback_model + assert guardrail in data["metadata"]["guardrails"] + assert data["messages"] == [{"role": "user", "content": "my ssn is [REDACTED-SSN]"}] + assert rig[3] == [primary_model, primary_model, fallback_model] + class _RecordingSuccessLogger(CustomLogger): def __init__(self): From 318b027782b0b6ab1aec7103a619d308661894e6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 08:26:39 +0000 Subject: [PATCH 041/102] fix(otel): keep caller traceparent and tracestate on pass-through relays Pass-through requests inject the proxy span into upstream headers since #40669, which replaced an explicit x-pass-traceparent with an unrelated trace and dropped its x-pass-tracestate. Keep the caller's context when the carrier already names a different trace, and keep the proxy child span for same-trace or missing headers. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 19 +++++- .../otel/test_otel_v2_components.py | 63 ++++++++++++++----- .../test_pass_through_endpoints.py | 59 +++++++++++++++++ 3 files changed, 125 insertions(+), 16 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 1282e654365..19243d64c64 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -325,12 +325,27 @@ def _outgoing_trace_context(parent_span: object) -> Context | None: return None +def _propagated_context(headers: Mapping[str, str], request_context: Context) -> Context: + """``request_context`` when it continues the trace ``headers`` already name, else the + caller's own context, so an explicit upstream ``traceparent`` (``x-pass-traceparent``) + is never swapped for an unrelated trace and its ``tracestate`` survives.""" + caller: Final = extract_traceparent(headers) + if caller is None: + return request_context + caller_span: Final = get_current_span(caller).get_span_context() + request_span: Final = get_current_span(request_context).get_span_context() + if not caller_span.is_valid or caller_span.trace_id == request_span.trace_id: + return request_context + return caller + + def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]: """``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span. Parent preference: ``parent_span`` (the request span auth stashed on the key), then the anchored request root span, then the ambient active span. Only trace context is - injected, never Baggage. Unchanged when no valid span exists anywhere. + injected, never Baggage. Unchanged when no valid span exists anywhere. A ``traceparent`` + already in ``headers`` from a different trace is forwarded as-is instead of replaced. """ context: Final = _outgoing_trace_context(parent_span) if context is None: @@ -338,7 +353,7 @@ def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS } - _PROPAGATOR.inject(carrier, context=context) + _PROPAGATOR.inject(carrier, context=_propagated_context(headers, context)) return carrier diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 72f6213e880..0c95049ce05 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -477,19 +477,21 @@ def _test_tracer(): return provider.get_tracer("test") +_CALLER_TRACEPARENT = "00-11111111111111111111111111111111-2222222222222222-01" + + def test_inject_trace_context_prefers_request_root_span(): def run(): tracer = _test_tracer() - with tracer.start_as_current_span("root") as root: + inbound = TraceContextTextMapPropagator().extract({"traceparent": _CALLER_TRACEPARENT}) + with tracer.start_as_current_span("root", context=inbound) as root: ctx_mod.set_request_root_span(root) - result = ctx_mod.inject_trace_context( - {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"} - ) + result = ctx_mod.inject_trace_context({"traceparent": _CALLER_TRACEPARENT}) propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) return result, root, propagated result, root, propagated = ContextVarContext().run(run) - assert result["traceparent"] != "00-11111111111111111111111111111111-2222222222222222-01" + assert result["traceparent"] != _CALLER_TRACEPARENT assert propagated.get_span_context().trace_id == root.get_span_context().trace_id assert propagated.get_span_context().span_id == root.get_span_context().span_id @@ -507,24 +509,57 @@ def test_inject_trace_context_uses_ambient_span_without_request_root(): assert propagated.get_span_context().span_id == ambient.get_span_context().span_id -def test_inject_trace_context_replaces_stale_trace_headers(): +def test_inject_trace_context_replaces_same_trace_headers_with_request_span(): def run(): tracer = _test_tracer() - with tracer.start_as_current_span("ambient") as ambient: - headers = { - "Traceparent": "00-" + "a" * 32 + "-" + "b" * 16 + "-01", - "Tracestate": "vendor=old", - "x-keep": "1", - } + headers = {"Traceparent": _CALLER_TRACEPARENT, "Tracestate": "vendor=caller", "x-keep": "1"} + inbound = TraceContextTextMapPropagator().extract({key.lower(): value for key, value in headers.items()}) + with tracer.start_as_current_span("ambient", context=inbound) as ambient: result = ctx_mod.inject_trace_context(headers) propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) return result, ambient, propagated result, ambient, propagated = ContextVarContext().run(run) assert sum(key.lower() == "traceparent" for key in result) == 1 - assert not any(key.lower() == "tracestate" for key in result) + assert sum(key.lower() == "tracestate" for key in result) == 1 assert result["x-keep"] == "1" - assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id + assert result["tracestate"] == "vendor=caller" + assert propagated.get_span_context().span_id == ambient.get_span_context().span_id + + +def test_inject_trace_context_keeps_caller_traceparent_from_another_trace(): + def run(): + tracer = _test_tracer() + parent = tracer.start_span("litellm_request") + with tracer.start_as_current_span("ambient") as ambient: + ctx_mod.set_request_root_span(ambient) + headers = {"Traceparent": _CALLER_TRACEPARENT, "Tracestate": "vendor=caller", "x-keep": "1"} + result = ctx_mod.inject_trace_context(headers, parent_span=parent) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return result, parent, propagated + + result, parent, propagated = ContextVarContext().run(run) + assert result["traceparent"] == _CALLER_TRACEPARENT + assert result["tracestate"] == "vendor=caller" + assert result["x-keep"] == "1" + assert sum(key.lower() == "traceparent" for key in result) == 1 + assert sum(key.lower() == "tracestate" for key in result) == 1 + assert propagated.get_span_context().trace_id != parent.get_span_context().trace_id + + +def test_inject_trace_context_replaces_malformed_caller_traceparent(): + def run(): + tracer = _test_tracer() + parent = tracer.start_span("litellm_request") + with tracer.start_as_current_span("ambient"): + headers = {"traceparent": "not-a-traceparent", "tracestate": "vendor=caller"} + result = ctx_mod.inject_trace_context(headers, parent_span=parent) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return result, parent, propagated + + result, parent, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().span_id == parent.get_span_context().span_id + assert "tracestate" not in result def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index e3e7ad618e0..df4842a7887 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4328,6 +4328,65 @@ async def test_pass_through_request_propagates_active_trace_context(span_source: assert propagated.get_span_context().span_id == span.get_span_context().span_id +async def _relay_with_trace_headers(inbound_headers: dict[str, str], forward_headers: bool): + from opentelemetry.sdk.trace import TracerProvider + + captured: dict[str, httpx.Headers] = {} + + def transport_handler(upstream_request: httpx.Request) -> httpx.Response: + captured["headers"] = upstream_request.headers + return httpx.Response(200, json={"ok": True}, request=upstream_request) + + fake_client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(transport_handler), timeout=None) + tracer = TracerProvider().get_tracer("test") + try: + with ExitStack() as stack: + _enter_relay_logging_mocks(stack, {}) + span = tracer.start_span("litellm_request") + stack.callback(span.end) + request = _relay_client_request(method="POST") + request.headers = Headers(inbound_headers) + response = await pass_through_request( + request=request, + target="http://internal-api.test/v1/generate", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=span), + forward_headers=forward_headers, + ) + finally: + cleanup() + await fake_client.aclose() + assert response.status_code == 200 + return captured["headers"], span + + +@pytest.mark.asyncio +@pytest.mark.parametrize("forward_headers", [False, True]) +async def test_pass_through_request_keeps_x_pass_trace_headers_when_otel_span_is_active(forward_headers: bool): + caller_traceparent = "00-11111111111111111111111111111111-2222222222222222-01" + + upstream_headers, span = await _relay_with_trace_headers( + {"x-pass-traceparent": caller_traceparent, "x-pass-tracestate": "vendor=caller"}, + forward_headers=forward_headers, + ) + + assert upstream_headers["traceparent"] == caller_traceparent + assert upstream_headers["tracestate"] == "vendor=caller" + assert format(span.get_span_context().trace_id, "032x") not in upstream_headers["traceparent"] + + +@pytest.mark.asyncio +async def test_pass_through_request_without_caller_trace_headers_still_propagates_proxy_span(): + from opentelemetry.trace import get_current_span + from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + + upstream_headers, span = await _relay_with_trace_headers({"x-pass-anthropic-beta": "beta-1"}, forward_headers=False) + + propagated = get_current_span(TraceContextTextMapPropagator().extract(upstream_headers)) + assert propagated.get_span_context().span_id == span.get_span_context().span_id + assert upstream_headers["anthropic-beta"] == "beta-1" + + @pytest.mark.asyncio async def test_pass_through_request_relays_non_json_body_without_buffering(): """ From b346eefd1f230b38f32e37ecb7a4ee8f73effb8b Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 08:33:46 +0000 Subject: [PATCH 042/102] fix(proxy): rerun requested model guardrail merge on fallback instead of carrying a raw list Structured guardrail entries (dicts) are unhashable, so merging a carried list through _merge_guardrails_with_existing raised TypeError on the fallback path. Resolve the requested model alias through _check_and_merge_model_level_guardrails instead and add a regression test for structured request guardrails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 31 +++++++------------ litellm/proxy/utils.py | 8 +++-- .../proxy/test_common_request_processing.py | 20 ++++++++++++ 3 files changed, 36 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 607ce3441e9..174b9ceff93 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -94,11 +94,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.route_llm_request import route_request -from litellm.proxy.utils import ( - ProxyLogging, - _check_and_merge_model_level_guardrails, - _merge_guardrails_with_existing, -) +from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_utils.common_utils import resolve_model_group_alias @@ -1938,7 +1934,7 @@ class ProxyBaseLLMRequestProcessing: user_api_base: str | None = None, model: str | None = None, llm_router: Router | None = None, - requested_model_guardrails: list | None = None, + rate_limited_model: str | None = None, ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks @@ -2102,12 +2098,15 @@ class ProxyBaseLLMRequestProcessing: # model_info when allow_client_pricing_override is set, so a caller # could otherwise spoof an unguarded model_info.id while requesting # a guarded alias and bypass guardrails (veria-ai HIGH on #29654). + merged_for_requested: Final = ( + self.data + if rate_limited_model is None + else _check_and_merge_model_level_guardrails( + data=self.data, llm_router=llm_router, trust_client_model_info=False, model_alias=rate_limited_model + ) + ) self.data = _check_and_merge_model_level_guardrails( - data=( - self.data - if requested_model_guardrails is None - else _merge_guardrails_with_existing(data=self.data, model_level_guardrails=requested_model_guardrails) - ), + data=merged_for_requested, llm_router=llm_router, trust_client_model_info=False, ) @@ -2217,8 +2216,6 @@ class ProxyBaseLLMRequestProcessing: original_model, fallback_models, ) - requested_model_guardrails: Final = self._request_guardrails(rate_limited_data) - try: for fallback_model in fallback_models: if fallback_model == original_model: @@ -2241,7 +2238,7 @@ class ProxyBaseLLMRequestProcessing: model=fallback_model, route_type=route_type, llm_router=llm_router, - requested_model_guardrails=requested_model_guardrails, + rate_limited_model=original_model, ) except ProxyRateLimitError: continue @@ -2252,12 +2249,6 @@ class ProxyBaseLLMRequestProcessing: self.data = rate_limited_data raise original_exc - @staticmethod - def _request_guardrails(data: dict) -> list | None: - metadata: Final = data.get("metadata") - guardrails: Final = metadata.get("guardrails") if isinstance(metadata, dict) else None - return guardrails if isinstance(guardrails, list) else None - @staticmethod def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None: key_router_settings: Final = user_api_key_dict.router_settings diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 950ac5e9906..ad0c0a82d49 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7497,6 +7497,7 @@ def _check_and_merge_model_level_guardrails( data: dict, llm_router: Router | None, trust_client_model_info: bool = True, + model_alias: str | None = None, ) -> dict: """ Check if the model has guardrails defined and merge them with existing guardrails in the request data. @@ -7504,6 +7505,7 @@ def _check_and_merge_model_level_guardrails( Args: data: The request data dict llm_router: The LLM router instance to get deployment info from + model_alias: Resolve guardrails for this model group instead of data["model"] trust_client_model_info: If False, ignore metadata.model_info.id and resolve guardrails by alias-union only. Set to False on the pre_call path because add_litellm_data_to_request preserves @@ -7548,13 +7550,13 @@ def _check_and_merge_model_level_guardrails( # set on ANY eligible deployment still fires (#29652; addresses # veria-ai HIGH on the single-deployment fallback that would skip # non-first deployments). - model_alias: Final = data.get("model") - if not isinstance(model_alias, str) or not model_alias: + alias: Final = model_alias if model_alias is not None else data.get("model") + if not isinstance(alias, str) or not alias: return data # Pass team_id so team-scoped public model names resolve the same way # route_request resolves them; otherwise team-scoped deployments are # invisible to this lookup and their guardrails are silently dropped. - deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] + deployments: Final = llm_router.get_model_list(model_name=alias, team_id=team_id) or [] seen: Final[set] = set() union: Final[list] = [] for dep in deployments: diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f440711ac29..135ce019813 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6970,6 +6970,26 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert data["messages"] == [{"role": "user", "content": "my ssn is [REDACTED-SSN]"}] assert rig[3] == [primary_model, primary_model, fallback_model] + @pytest.mark.asyncio + async def test_fallback_keeps_structured_request_guardrails(self, monkeypatch: pytest.MonkeyPatch): + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + structured_guardrail = {"pii-guard": {"extra_body": {"threshold": 0.5}}} + key = self._otel_key(model_rpm_limit={primary_model: 1}) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request = { + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "guardrails": [structured_guardrail], + } + + await self._pre_call(dict(request), key, rig) + _, (data, _) = await self._pre_call(dict(request), key, rig) + + assert data["model"] == fallback_model + assert data["metadata"]["guardrails"] == [structured_guardrail] + assert rig[3] == [primary_model, primary_model, fallback_model] + class _RecordingSuccessLogger(CustomLogger): def __init__(self): From dba9ff801f6b3474e7861b7c066396ebe19e3cf3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:18:49 +0000 Subject: [PATCH 043/102] fix(proxy): reset sibling tpm/rpm counters when shared window rolls over Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../hooks/parallel_request_limiter_v3.py | 39 +++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 69 +++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 89826694bb6..57b7d4f249d 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -91,10 +91,16 @@ else: _REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object]) +def _sibling_counter_keys(window_key: str) -> tuple[str, str]: + prefix: Final = window_key.removesuffix(":window") + return f"{prefix}:requests", f"{prefix}:tokens" + + BATCH_RATE_LIMITER_SCRIPT: Final = """ local results = {} local now = tonumber(ARGV[1]) local window_size = tonumber(ARGV[2]) +local reset_windows = {} -- Process each window/counter pair for i = 1, #KEYS, 2 do @@ -106,6 +112,11 @@ for i = 1, #KEYS, 2 do local window_start = redis.call('GET', window_key) if not window_start or (now - tonumber(window_start)) >= window_size then -- Reset window and counter + if not reset_windows[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset_windows[window_key] = true + end redis.call('SET', window_key, tostring(now)) redis.call('SET', counter_key, increment_value) redis.call('EXPIRE', window_key, window_size) @@ -151,6 +162,7 @@ CHECK_AND_INCREMENT_BY_N_SCRIPT: Final = """ local time_reply = redis.call('TIME') local now = tonumber(time_reply[1]) local descriptor_count = #KEYS / 2 +local reset_windows = {} -- Pass 1: read state, validate. Abort without writing if any over limit. local descriptor_state = {} @@ -201,6 +213,11 @@ for i = 1, descriptor_count do if window_expired then active_window_start = now + if not reset_windows[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset_windows[window_key] = true + end redis.call('SET', window_key, tostring(now)) redis.call('SET', counter_key, increment) redis.call('EXPIRE', window_key, window_size) @@ -1019,6 +1036,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): This follows the same logic as the Redis Lua script but uses async cache operations. """ results: Final[list[CacheCounterValue | None]] = [] + reset_windows: Final[set[str]] = set() # mutable-ok: tracks windows reset during this call # Process each window/counter pair for i in range(0, len(keys), 2): @@ -1036,6 +1054,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Check if window exists and is valid if window_start is None or (now_int - int(window_start)) >= window_size: # Reset window and counter + if window_key not in reset_windows: + for sibling_counter_key in _sibling_counter_keys(window_key): + await self.internal_usage_cache.async_set_cache( + key=sibling_counter_key, + value=0, + ttl=window_size, + litellm_parent_otel_span=None, + local_only=True, + ) + reset_windows.add(window_key) await self.internal_usage_cache.async_set_cache( key=window_key, value=str(now_int), @@ -2049,9 +2077,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Pass 2: apply increments. statuses: Final[list[RateLimitStatus]] = [] + reset_windows: Final[set[str]] = set() # mutable-ok: tracks windows reset during this call for meta, state in zip(per_counter_meta, descriptor_state): new_counter = meta["increment"] if state["window_expired"] else state["current"] + meta["increment"] if state["window_expired"]: + if meta["window_key"] not in reset_windows: + for sibling_counter_key in _sibling_counter_keys(meta["window_key"]): + await self.internal_usage_cache.async_set_cache( + key=sibling_counter_key, + value=0, + ttl=meta["window_size"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + reset_windows.add(meta["window_key"]) await self.internal_usage_cache.async_set_cache( key=meta["window_key"], value=str(now_int), diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 85023b94207..0b7710eda70 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5602,6 +5602,75 @@ async def _reserved_tokens_for( return int(await local_cache.async_get_cache(key=tokens_key) or 0) +@pytest.mark.asyncio +async def test_tpm_reservation_resets_sibling_tokens_with_request_window(monkeypatch): + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "true") + time_controller = TimeController() + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=time_controller.now, + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-window-reset-siblings"), + tpm_limit=1000, + rpm_limit=1000, + ) + + async def request(call_id): + data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 200, + "litellm_call_id": call_id, + "metadata": { + "user_api_key": user_api_key_dict.api_key, + "user_api_key_user_id": user_api_key_dict.user_id, + }, + } + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type="completion", + ) + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": call_id, + "litellm_params": { + "metadata": { + "user_api_key": user_api_key_dict.api_key, + "user_api_key_user_id": user_api_key_dict.user_id, + "model_group": "gpt-4o", + } + }, + "standard_logging_object": { + "metadata": { + "user_api_key_hash": user_api_key_dict.api_key, + "user_api_key_user_id": user_api_key_dict.user_id, + } + }, + }, + response_obj=ModelResponse( + model="gpt-4o", + usage=Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300), + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + tokens_key = handler.create_rate_limit_keys( + key="api_key", value=user_api_key_dict.api_key, rate_limit_type="tokens" + ) + for index in range(3): + await request(f"call-{index}") + assert await local_cache.async_get_cache(key=tokens_key) == (index + 1) * 300 + + time_controller.advance(61) + await request("call-after-window-reset") + assert await local_cache.async_get_cache(key=tokens_key) == 300 + + @pytest.mark.asyncio @pytest.mark.parametrize( "key_metadata, team_metadata, expected_output_estimate, tier", From 3219a875d0dc350ba129b881fa1d9dc003ba5b17 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 18:20:18 +0000 Subject: [PATCH 044/102] fix(proxy): serialize in-memory window rollover and reset siblings once per expired window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../hooks/parallel_request_limiter_v3.py | 61 +++++++------- .../hooks/test_parallel_request_limiter_v3.py | 82 +++++++++++++++++++ 2 files changed, 115 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 57b7d4f249d..cdec5922fff 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -100,7 +100,6 @@ BATCH_RATE_LIMITER_SCRIPT: Final = """ local results = {} local now = tonumber(ARGV[1]) local window_size = tonumber(ARGV[2]) -local reset_windows = {} -- Process each window/counter pair for i = 1, #KEYS, 2 do @@ -112,11 +111,8 @@ for i = 1, #KEYS, 2 do local window_start = redis.call('GET', window_key) if not window_start or (now - tonumber(window_start)) >= window_size then -- Reset window and counter - if not reset_windows[window_key] then - local prefix = string.sub(window_key, 1, -(#':window') - 1) - redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') - reset_windows[window_key] = true - end + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') redis.call('SET', window_key, tostring(now)) redis.call('SET', counter_key, increment_value) redis.call('EXPIRE', window_key, window_size) @@ -1035,8 +1031,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Implement sliding window rate limiting logic using in-memory cache operations. This follows the same logic as the Redis Lua script but uses async cache operations. """ + async with self._check_and_increment_lock: + return await self._in_memory_cache_sliding_window(keys=keys, now_int=now_int, window_size=window_size) + + async def _in_memory_cache_sliding_window( + self, + keys: list[str], + now_int: int, + window_size: int, + ) -> CacheCounterValues: results: Final[list[CacheCounterValue | None]] = [] - reset_windows: Final[set[str]] = set() # mutable-ok: tracks windows reset during this call # Process each window/counter pair for i in range(0, len(keys), 2): @@ -1054,16 +1058,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Check if window exists and is valid if window_start is None or (now_int - int(window_start)) >= window_size: # Reset window and counter - if window_key not in reset_windows: - for sibling_counter_key in _sibling_counter_keys(window_key): - await self.internal_usage_cache.async_set_cache( - key=sibling_counter_key, - value=0, - ttl=window_size, - litellm_parent_otel_span=None, - local_only=True, - ) - reset_windows.add(window_key) + for sibling_counter_key in _sibling_counter_keys(window_key): + await self.internal_usage_cache.async_set_cache( + key=sibling_counter_key, + value=0, + ttl=window_size, + litellm_parent_otel_span=None, + local_only=True, + ) await self.internal_usage_cache.async_set_cache( key=window_key, value=str(now_int), @@ -2076,21 +2078,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) # Pass 2: apply increments. + expired_windows: Final[Mapping[str, int]] = { + meta["window_key"]: meta["window_size"] + for meta, state in zip(per_counter_meta, descriptor_state) + if state["window_expired"] + } + for window_key, window_size in expired_windows.items(): + for sibling_counter_key in _sibling_counter_keys(window_key): + await self.internal_usage_cache.async_set_cache( + key=sibling_counter_key, + value=0, + ttl=window_size, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) statuses: Final[list[RateLimitStatus]] = [] - reset_windows: Final[set[str]] = set() # mutable-ok: tracks windows reset during this call for meta, state in zip(per_counter_meta, descriptor_state): new_counter = meta["increment"] if state["window_expired"] else state["current"] + meta["increment"] if state["window_expired"]: - if meta["window_key"] not in reset_windows: - for sibling_counter_key in _sibling_counter_keys(meta["window_key"]): - await self.internal_usage_cache.async_set_cache( - key=sibling_counter_key, - value=0, - ttl=meta["window_size"], - litellm_parent_otel_span=parent_otel_span, - local_only=True, - ) - reset_windows.add(meta["window_key"]) await self.internal_usage_cache.async_set_cache( key=meta["window_key"], value=str(now_int), diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 0b7710eda70..5889b1b513f 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -18,11 +18,13 @@ from fastapi import HTTPException import litellm from litellm import Router from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PARALLEL_REQUEST_SLOT_TTL_SECONDS, ParallelSlotAcquisition, + RateLimitDescriptor, RequestRateLimiterStash, _request_stash, get_or_create_request_stash, @@ -5671,6 +5673,86 @@ async def test_tpm_reservation_resets_sibling_tokens_with_request_window(monkeyp assert await local_cache.async_get_cache(key=tokens_key) == 300 +@pytest.mark.asyncio +async def test_atomic_tpm_reservation_rollover_resets_sibling_requests_counter(): + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + window_size = 60 + now_int = int(time.time()) + window_key = "{api_key:atomic-rollover}:window" + requests_key = handler.create_rate_limit_keys("api_key", "atomic-rollover", "requests") + tokens_key = handler.create_rate_limit_keys("api_key", "atomic-rollover", "tokens") + for key, value in ((window_key, str(now_int - window_size - 1)), (requests_key, 3), (tokens_key, 900)): + await local_cache.async_set_cache(key=key, value=value, ttl=window_size) + + tpm_pass = await handler.atomic_check_and_increment_by_n( + descriptors=[ + RateLimitDescriptor( + key="api_key", + value="atomic-rollover", + rate_limit={"tokens_per_unit": 1000, "window_size": window_size}, + ) + ], + increments=[{"tokens": 200}], + ) + assert tpm_pass["overall_code"] == "OK" + assert await local_cache.async_get_cache(key=tokens_key) == 200 + + rpm_pass = await handler.should_rate_limit( + descriptors=[ + RateLimitDescriptor( + key="api_key", + value="atomic-rollover", + rate_limit={"requests_per_unit": 5, "window_size": window_size}, + ) + ], + skip_tpm_check=True, + ) + assert rpm_pass["overall_code"] == "OK" + assert [status["limit_remaining"] for status in rpm_pass["statuses"]] == [4] + assert await local_cache.async_get_cache(key=requests_key) == 1 + + +class _YieldingInMemoryCache(InMemoryCache): + async def async_get_cache(self, key: str, **kwargs: object) -> object: + value = await super().async_get_cache(key, **kwargs) + await asyncio.sleep(0) + return value + + +@pytest.mark.asyncio +async def test_window_rollover_reset_does_not_erase_concurrent_sibling_increment(): + local_cache = DualCache(in_memory_cache=_YieldingInMemoryCache()) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + window_size = 60 + now_int = int(time.time()) + window_key = "{api_key:concurrent-rollover}:window" + requests_key = handler.create_rate_limit_keys("api_key", "concurrent-rollover", "requests") + tokens_key = handler.create_rate_limit_keys("api_key", "concurrent-rollover", "tokens") + for key, value in ((window_key, str(now_int - window_size - 1)), (requests_key, 3), (tokens_key, 900)): + await local_cache.async_set_cache(key=key, value=value, ttl=window_size) + + tpm_descriptor = RateLimitDescriptor( + key="api_key", + value="concurrent-rollover", + rate_limit={"tokens_per_unit": 1000, "window_size": window_size}, + ) + rpm_pass, tpm_pass = await asyncio.gather( + handler.in_memory_cache_sliding_window( + keys=[window_key, requests_key], now_int=now_int, window_size=window_size + ), + handler.atomic_check_and_increment_by_n( + descriptors=[tpm_descriptor], + increments=[{"requests": 0, "tokens": 200}], + ), + ) + + assert rpm_pass == [str(now_int), 1] + assert tpm_pass["overall_code"] == "OK" + assert await local_cache.async_get_cache(key=requests_key) == 1 + assert await local_cache.async_get_cache(key=tokens_key) == 200 + + @pytest.mark.asyncio @pytest.mark.parametrize( "key_metadata, team_metadata, expected_output_estimate, tier", From cf08440557bd0a377f1ce082497d037294c6aca0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:42:52 -0700 Subject: [PATCH 045/102] fix(azure-ai): price FLUX.2 Flex by its resolved cost map key --- .../azure_ai/image_generation/cost_calculator.py | 5 +---- .../test_azure_ai_flux2_image_generation.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 086293f26c0..35d0f4fb6c3 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -42,9 +42,6 @@ def cost_calculator( if input_cost_per_pixel: from litellm.cost_calculator import default_image_cost_calculator - cost_model: Final = ( - model if model.startswith(f"{litellm.LlmProviders.AZURE_AI.value}/") else f"azure_ai/{model}" - ) width: Final = optional_params.get("width") if optional_params else None height: Final = optional_params.get("height") if optional_params else None pixel_size: Final = ( @@ -53,7 +50,7 @@ def cost_calculator( else size or image_response.size ) return default_image_cost_calculator( - model=cost_model, + model=_model_info["key"], custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, size=pixel_size, n=num_images, diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index c1d46b7e919..388cca9f054 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -172,6 +172,19 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi ) == pytest.approx(5e-08 * 2048 * 1024 * 2) +def test_flux2_flex_cost_accepts_lowercase_model_spelling(): + response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")]) + + cost: Final = litellm.completion_cost( + model="azure_ai/flux.2-flex", + completion_response=response, + optional_params={"width": 1536, "height": 1024, "num_images": 2}, + call_type="image_generation", + ) + + assert cost == pytest.approx(5e-08 * 1536 * 1024 * 2) + + def test_flux2_response_preserves_mapped_dimensions(): config = AzureFoundryFluxImageGenerationConfig() params = config.map_openai_params( From c1e39810ecbcc3def67502932273b5e7a0bc94a8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:58:04 -0700 Subject: [PATCH 046/102] fix(mistral): send reasoning_effort as a level the model accepts Declare the live-verified reasoning_effort_levels on the Mistral cost-map entries and round an undeclared request to the nearest declared level (up to the weakest level at least as strong, down to the strongest when the request exceeds the ceiling). Codex's default medium no longer 400s on mistral-medium-latest, mistral-small-latest, or the vibe-cli family; an entry that declares nothing keeps forwarding the value verbatim --- litellm/llms/mistral/chat/transformation.py | 19 ++++- ...odel_prices_and_context_window_backup.json | 77 +++++++++++++++++++ .../reasoning_effort_capability.py | 17 ++++ model_prices_and_context_window.json | 77 +++++++++++++++++++ .../test_mistral_chat_transformation.py | 44 +++++++++-- .../test_reasoning_effort_capability.py | 20 +++++ 6 files changed, 246 insertions(+), 8 deletions(-) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 6316128e6fa..8cbe4291054 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast, get_type_hints, ove import httpx +from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, @@ -20,6 +21,10 @@ from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) +from litellm.router_utils.reasoning_effort_capability import ( + declared_reasoning_efforts_for_model, + nearest_declared_reasoning_effort, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.mistral import MistralThinkingBlock, MistralToolCallMessage from litellm.types.llms.openai import AllMessageValues @@ -30,6 +35,18 @@ if TYPE_CHECKING: import tiktoken +def _accepted_reasoning_effort(model: str, requested: str) -> str: + declared: Final = declared_reasoning_efforts_for_model(model, "mistral") + if declared is None: + return requested + accepted: Final = nearest_declared_reasoning_effort(requested, declared) + if accepted != requested: + verbose_logger.debug( + "mistral: %s takes reasoning_effort %s, sending %s in place of %s", model, declared, accepted, requested + ) + return accepted + + class MistralConfig(OpenAIGPTConfig): """ Reference: https://docs.mistral.ai/api/ @@ -170,7 +187,7 @@ class MistralConfig(OpenAIGPTConfig): if param == "response_format": optional_params["response_format"] = value if param == "reasoning_effort" and "magistral" not in model.lower(): - optional_params["reasoning_effort"] = value + optional_params["reasoning_effort"] = _accepted_reasoning_effort(model, value) if param in ("reasoning_effort", "thinking") and "magistral" in model.lower(): # Flag that we need to add reasoning system prompt optional_params["_add_reasoning_prompt"] = True diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0826f05cb11..f8707f9e403 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -36993,6 +36993,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37072,6 +37076,15 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-2", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37089,6 +37102,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-3", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37106,6 +37124,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-3", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37123,6 +37146,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-3", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37140,6 +37168,15 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-2", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37435,6 +37472,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37495,6 +37536,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37512,6 +37557,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37545,6 +37594,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37575,6 +37628,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -59375,6 +59432,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -62700,6 +62761,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -62717,6 +62782,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -62734,6 +62803,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -62751,6 +62824,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index 7b145c15a07..f880c846500 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -103,6 +103,23 @@ def declared_reasoning_efforts_for_model(model: str, custom_llm_provider: str) - return declared_reasoning_efforts(entry) +REASONING_EFFORT_STRENGTH_ORDER: Final = ("none", "minimal", "low", "medium", "high", "xhigh", "max") +_STRENGTH_RANK: Final = MappingProxyType({effort: rank for rank, effort in enumerate(REASONING_EFFORT_STRENGTH_ORDER)}) + + +def nearest_declared_reasoning_effort(requested: str, declared: Sequence[str]) -> str: + """Rounds a request up to the weakest declared level at least as strong as it, and down to the + strongest declared level when it asks for more than the model has, so the caller gets no less + reasoning than it asked for instead of a rejected call. A level outside the strength order is + returned as is for upstream to judge.""" + ranked: Final = sorted( + (effort for effort in declared if effort in _STRENGTH_RANK), key=lambda effort: _STRENGTH_RANK[effort] + ) + if requested in ranked or requested not in _STRENGTH_RANK or not ranked: + return requested + return next((effort for effort in ranked if _STRENGTH_RANK[effort] >= _STRENGTH_RANK[requested]), ranked[-1]) + + def _supports_none_reasoning_effort(model_info: Mapping[str, object], flag: object) -> bool: """Opt-in only where a request path refuses the level. AzureOpenAIGPT5Config raises UnsupportedParamsError on reasoning_effort='none' without an explicit true, and it is selected diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0826f05cb11..f8707f9e403 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -36993,6 +36993,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37072,6 +37076,15 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-2", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37089,6 +37102,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-3", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37106,6 +37124,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-3", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37123,6 +37146,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-3", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37140,6 +37168,15 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.4e-06, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ], "source": "https://docs.mistral.ai/models/zai-glm-5-2", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37435,6 +37472,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37495,6 +37536,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37512,6 +37557,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37545,6 +37594,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -37575,6 +37628,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -59375,6 +59432,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -62700,6 +62761,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -62717,6 +62782,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -62734,6 +62803,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 7.5e-06, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -62751,6 +62824,10 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 6e-07, + "reasoning_effort_levels": [ + "none", + "high" + ], "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py index 38639f23050..a68ba9f570c 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py @@ -63,7 +63,7 @@ class TestMistralReasoningSupport: assert "reasoning_effort" not in supported_params_normal assert "thinking" not in supported_params_normal - def test_map_openai_params_reasoning_effort(self): + def test_map_openai_params_reasoning_effort(self, local_model_cost_map): """Test that reasoning_effort parameter is properly mapped for magistral models.""" mistral_config = MistralConfig() @@ -87,21 +87,51 @@ class TestMistralReasoningSupport: ) assert "_add_reasoning_prompt" not in result_normal - assert result_normal["reasoning_effort"] == "low" + assert result_normal["reasoning_effort"] == "high" @pytest.mark.parametrize( - ("model", "reasoning_effort"), - [("mistral-medium-latest", "high"), ("zai-glm-5-2", "xhigh")], + ("model", "requested", "sent"), + [ + ("mistral-medium-latest", "high", "high"), + ("mistral-medium-latest", "none", "none"), + ("mistral-medium-latest", "low", "high"), + ("mistral-medium-latest", "medium", "high"), + ("mistral-medium-latest", "xhigh", "high"), + ("mistral-small-latest", "medium", "high"), + ("mistral-vibe-cli-latest", "medium", "high"), + ("zai-glm-5", "none", "low"), + ("zai-glm-5", "medium", "high"), + ("zai-glm-5", "xhigh", "max"), + ("zai-glm-5-2", "medium", "medium"), + ("zai-glm-5-2", "xhigh", "xhigh"), + ], ) - def test_reasoning_effort_forwarded_verbatim_for_reasoning_models(self, model, reasoning_effort): + def test_reasoning_effort_is_sent_as_a_level_the_model_accepts(self, local_model_cost_map, model, requested, sent): import litellm optional_params = litellm.get_optional_params( model=model, custom_llm_provider="mistral", - reasoning_effort=reasoning_effort, + reasoning_effort=requested, ) - assert optional_params["reasoning_effort"] == reasoning_effort + assert optional_params["reasoning_effort"] == sent + + def test_reasoning_effort_is_forwarded_verbatim_when_the_map_declares_no_levels( + self, local_model_cost_map, monkeypatch + ): + import litellm + + monkeypatch.setitem( + litellm.model_cost, + "mistral/undeclared-reasoner", + {"litellm_provider": "mistral", "mode": "chat", "supports_reasoning": True}, + ) + optional_params = litellm.get_optional_params( + model="undeclared-reasoner", + custom_llm_provider="mistral", + reasoning_effort="medium", + ) + assert optional_params["reasoning_effort"] == "medium" def test_reasoning_effort_stays_unsupported_for_non_reasoning_models(self): import litellm diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index ccd6766b13a..ca4a3517bfc 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -4,6 +4,7 @@ import litellm from litellm.router_utils.reasoning_effort_capability import ( deployment_is_catalog_mapped, intersect_supported_reasoning_efforts, + nearest_declared_reasoning_effort, resolve_supported_reasoning_efforts, ) @@ -415,3 +416,22 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: "high", "xhigh", ) + + +class TestNearestDeclaredReasoningEffort: + def test_a_declared_level_is_kept(self): + assert nearest_declared_reasoning_effort("high", ("none", "high")) == "high" + assert nearest_declared_reasoning_effort("none", ("none", "high")) == "none" + + def test_an_undeclared_level_rounds_up_to_the_next_declared_one(self): + assert nearest_declared_reasoning_effort("medium", ("none", "high")) == "high" + assert nearest_declared_reasoning_effort("none", ("low", "high", "max")) == "low" + assert nearest_declared_reasoning_effort("xhigh", ("low", "high", "max")) == "max" + + def test_a_level_above_the_ceiling_takes_the_strongest_declared_one(self): + assert nearest_declared_reasoning_effort("max", ("none", "high")) == "high" + assert nearest_declared_reasoning_effort("xhigh", ("none", "low", "medium", "high")) == "high" + + def test_a_level_outside_the_strength_order_is_left_for_upstream(self): + assert nearest_declared_reasoning_effort("turbo", ("none", "high")) == "turbo" + assert nearest_declared_reasoning_effort("medium", ()) == "medium" From 84f7adec2e4d07ec00cea03b75d0509cde817418 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 19:06:24 +0000 Subject: [PATCH 047/102] fix(passthrough): match deepgram listen routes served under a path prefix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../deepgram_listen_passthrough_logging_handler.py | 2 +- .../test_deepgram_listen_passthrough_logging_handler.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py index fc0af400477..c0953ae7b85 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py @@ -57,7 +57,7 @@ class DeepgramListenPassthroughLoggingHandler: @staticmethod def is_deepgram_listen_route(url_route: str) -> bool: path: Final = urlparse(url_route).path - return path.startswith("/deepgram/") and path.endswith(DEEPGRAM_LISTEN_ROUTE_SUFFIX) + return "/deepgram/" in path and path.endswith(DEEPGRAM_LISTEN_ROUTE_SUFFIX) def deepgram_listen_passthrough_handler( self, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py index a6133a7cf99..2ae56709b11 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py @@ -42,6 +42,7 @@ def _metadata(duration: object, channels: int = 1) -> dict[str, object]: ("/deepgram/v1/listen", True), ("/deepgram/listen", True), ("/deepgram/v1/listen?model=nova-3", True), + ("/litellm/deepgram/v1/listen", True), ("/deepgram/v1/speak", False), ("/deepgram/v1/listen/extra", False), ("/openai/v1/realtime", False), From aad89de4b4cb64aa8c921dd5315c236d46c1d7f0 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 19:13:56 +0000 Subject: [PATCH 048/102] build(deps): bump anyio to 4.14.2 in uv.lock to clear GHSA-5p39-cfhj-2xmp and GHSA-82r6-8w77-94w6 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- uv.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index a5e60c68515..04c31bea7d3 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-14T23:55:55.024292355Z" +exclude-newer = "2026-09-15T19:13:10.278189214Z" exclude-newer-span = "P3D" [manifest] @@ -315,16 +315,16 @@ vertex = [ [[package]] name = "anyio" -version = "4.13.0" +version = "4.14.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, { name = "idna" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } +sdist = { url = "https://files.pythonhosted.org/packages/61/cc/a381afa6efea9f496eff839d4a6a1aed3bfafc7b3ab4b0d1b243a12573dd/anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f", size = 260176, upload-time = "2026-07-12T20:29:07.082Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, + { url = "https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl", hash = "sha256:9f505dda5ac9f0c8309b5e8bd445a8c2bf7246f3ce950121e45ea15bc41d1494", size = 125813, upload-time = "2026-07-12T20:29:05.763Z" }, ] [[package]] From b5ad17e8482f924e1fafb71563d2e4819f105ddb Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 19:11:39 +0000 Subject: [PATCH 049/102] test(docs): read only the first column of the router_settings reference table Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/documentation_tests/test_router_settings.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/documentation_tests/test_router_settings.py b/tests/documentation_tests/test_router_settings.py index 75032f80dfa..7e3d0c07459 100644 --- a/tests/documentation_tests/test_router_settings.py +++ b/tests/documentation_tests/test_router_settings.py @@ -51,9 +51,7 @@ try: if general_settings_section: # Extract the table rows, which contain the documented keys table_content = general_settings_section.group(1) - doc_key_pattern = re.compile( - r"\|\s*([^\|]+?)\s*\|" - ) # Capture the key from each row of the table + doc_key_pattern = re.compile(r"^\|\s*([^\|]+?)\s*\|", re.MULTILINE) documented_keys.update(doc_key_pattern.findall(table_content)) except Exception as e: raise Exception( From b4dc081c275da7b631faa9cf59782a3cb5452141 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:22:34 -0700 Subject: [PATCH 050/102] fix(mistral): never round reasoning_effort none onto the strength ladder --- litellm/router_utils/reasoning_effort_capability.py | 8 +++++--- .../llms/mistral/test_mistral_chat_transformation.py | 3 ++- .../router_utils/test_reasoning_effort_capability.py | 6 +++++- 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index f880c846500..1d7656f253e 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -103,15 +103,17 @@ def declared_reasoning_efforts_for_model(model: str, custom_llm_provider: str) - return declared_reasoning_efforts(entry) -REASONING_EFFORT_STRENGTH_ORDER: Final = ("none", "minimal", "low", "medium", "high", "xhigh", "max") +REASONING_EFFORT_STRENGTH_ORDER: Final = ("minimal", "low", "medium", "high", "xhigh", "max") _STRENGTH_RANK: Final = MappingProxyType({effort: rank for rank, effort in enumerate(REASONING_EFFORT_STRENGTH_ORDER)}) def nearest_declared_reasoning_effort(requested: str, declared: Sequence[str]) -> str: """Rounds a request up to the weakest declared level at least as strong as it, and down to the strongest declared level when it asks for more than the model has, so the caller gets no less - reasoning than it asked for instead of a rejected call. A level outside the strength order is - returned as is for upstream to judge.""" + reasoning than it asked for instead of a rejected call. none is the off switch rather than a + strength, so it is never rounded onto the ladder and no level is rounded down to it: a caller + who turned reasoning off must not be billed for it, and a model that cannot turn it off says so + itself. A level outside the strength order is likewise returned as is for upstream to judge.""" ranked: Final = sorted( (effort for effort in declared if effort in _STRENGTH_RANK), key=lambda effort: _STRENGTH_RANK[effort] ) diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py index a68ba9f570c..8fb3b3c43df 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py @@ -99,7 +99,8 @@ class TestMistralReasoningSupport: ("mistral-medium-latest", "xhigh", "high"), ("mistral-small-latest", "medium", "high"), ("mistral-vibe-cli-latest", "medium", "high"), - ("zai-glm-5", "none", "low"), + ("zai-glm-5", "none", "none"), + ("zai-glm-5", "minimal", "low"), ("zai-glm-5", "medium", "high"), ("zai-glm-5", "xhigh", "max"), ("zai-glm-5-2", "medium", "medium"), diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index ca4a3517bfc..adee44aa8a3 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -425,9 +425,13 @@ class TestNearestDeclaredReasoningEffort: def test_an_undeclared_level_rounds_up_to_the_next_declared_one(self): assert nearest_declared_reasoning_effort("medium", ("none", "high")) == "high" - assert nearest_declared_reasoning_effort("none", ("low", "high", "max")) == "low" + assert nearest_declared_reasoning_effort("minimal", ("low", "high", "max")) == "low" assert nearest_declared_reasoning_effort("xhigh", ("low", "high", "max")) == "max" + def test_none_is_a_switch_that_is_never_rounded_in_either_direction(self): + assert nearest_declared_reasoning_effort("none", ("low", "high", "max")) == "none" + assert nearest_declared_reasoning_effort("medium", ("none",)) == "medium" + def test_a_level_above_the_ceiling_takes_the_strongest_declared_one(self): assert nearest_declared_reasoning_effort("max", ("none", "high")) == "high" assert nearest_declared_reasoning_effort("xhigh", ("none", "low", "medium", "high")) == "high" From c0b0ba20b23a83a5e382ee2b7de7739cf5dcaa10 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:41:40 -0700 Subject: [PATCH 051/102] fix(azure-ai): keep FLUX.2 tolerant of OpenAI-only image params --- .../image_edit/flux2_transformation.py | 17 +--------- .../image_generation/flux_transformation.py | 34 +++++++++++++++---- ...test_azure_ai_image_edit_transformation.py | 11 ++++++ .../test_azure_ai_flux2_image_generation.py | 34 +++++++++++++++---- 4 files changed, 67 insertions(+), 29 deletions(-) diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index aa8905e5601..f91a87ba0f4 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -31,22 +31,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ def get_supported_openai_params(self, model: str) -> list: - """ - FLUX 2 supports a subset of OpenAI image edit params - """ - return [ - "n", - "size", - "width", - "height", - "num_images", - "seed", - "safety_tolerance", - "output_format", - "aspect_ratio", - "guidance", - "steps", - ] + return AzureFoundryFluxImageGenerationConfig().get_supported_openai_params(model) def map_openai_params( self, diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 997205b1fc3..ac9ec24420b 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -2,9 +2,18 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final +from litellm.exceptions import BadRequestError, UnsupportedParamsError from litellm.llms.openai.image_generation import GPTImageGenerationConfig from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams +FLUX2_DROPPED_OPENAI_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = ( + "background", + "moderation", + "output_compression", + "quality", + "user", +) + class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): """Azure Foundry BFL API configuration for FLUX image generation.""" @@ -81,10 +90,13 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): "num_images", "guidance", "steps", + *FLUX2_DROPPED_OPENAI_PARAMS, ] @staticmethod - def _map_parameter(name: str, value: object) -> tuple[tuple[str, object], ...]: + def _map_parameter(name: str, value: object, model: str) -> tuple[tuple[str, object], ...]: + if name in FLUX2_DROPPED_OPENAI_PARAMS: + return () if isinstance(value, str): if name in ("n", "num_images", "width", "height", "steps", "seed", "safety_tolerance"): return (("num_images" if name == "n" else name, int(value)),) @@ -94,11 +106,17 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): return (("num_images", value),) if name != "size": return ((name, value),) + if str(value).lower() == "auto": + return () try: width, height = (int(dimension) for dimension in str(value).lower().split("x")) except (TypeError, ValueError): - raise ValueError(f"Invalid size format '{value}'. Expected 'WxH', for example '1024x1024'.") + raise BadRequestError( + message=f"Invalid size format '{value}'. Expected 'WxH', for example '1024x1024'.", + model=model, + llm_provider="azure_ai", + ) return (("width", width), ("height", height)) def map_openai_params( @@ -118,9 +136,13 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): supported_params: Final = self.get_supported_openai_params(model) unsupported_params: Final = tuple(name for name in non_default_params if name not in supported_params) if unsupported_params and not drop_params: - raise ValueError( - f"Parameters {unsupported_params} are not supported for model {model}. " - f"Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." + raise UnsupportedParamsError( + message=( + f"Parameters {unsupported_params} are not supported for model {model}. " + f"Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." + ), + model=model, + llm_provider="azure_ai", ) mapped_params: Final[Mapping[str, object]] = MappingProxyType( @@ -128,7 +150,7 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): mapped_name: mapped_value for name, value in non_default_params.items() if name in supported_params - for mapped_name, mapped_value in self._map_parameter(name, value) + for mapped_name, mapped_value in self._map_parameter(name, value, model) } ) return {**optional_params, **mapped_params} # mutable-ok: inherited config contract returns a dict diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index d74afa88a6b..39001c1795b 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -176,3 +176,14 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ ) assert response._hidden_params["response_cost"] == pytest.approx(5e-08 * 2048 * 1024 * 2) + + +def test_flux2_image_edit_accepts_and_drops_openai_only_parameters(): + optional_params: Final = ImageEditRequestUtils.get_optional_params_image_edit( + model="FLUX.2-pro", + image_edit_provider_config=AzureFoundryFlux2ImageEditConfig(), + image_edit_optional_params={"n": 1, "size": "auto", "quality": "high", "user": "end-user-1"}, + drop_params=False, + ) + + assert optional_params == {"num_images": 1} diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index 388cca9f054..512e98b4151 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -15,7 +15,7 @@ from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) from litellm.types.utils import ImageObject, ImageResponse -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params_image_gen @pytest.fixture(autouse=True) @@ -86,15 +86,35 @@ def test_flux2_flex_maps_openai_and_provider_parameters(): } -def test_flux2_flex_rejects_invalid_size(): - with pytest.raises(ValueError, match="Expected 'WxH'"): - AzureFoundryFluxImageGenerationConfig().map_openai_params( - non_default_params={"size": "large"}, - optional_params={}, +def test_flux2_flex_rejects_invalid_size_as_bad_request(): + with pytest.raises(litellm.BadRequestError, match="Expected 'WxH'") as raised: + get_optional_params_image_gen( model="FLUX.2-flex", - drop_params=False, + custom_llm_provider="azure_ai", + provider_config=AzureFoundryFluxImageGenerationConfig(), + size="large", ) + assert raised.value.status_code == 400 + + +@pytest.mark.parametrize("model", ("FLUX.2-pro", "FLUX.2-flex")) +def test_flux2_accepts_and_drops_openai_only_image_parameters(model: str): + optional_params: Final = get_optional_params_image_gen( + model=model, + custom_llm_provider="azure_ai", + provider_config=AzureFoundryFluxImageGenerationConfig(), + n=1, + size="auto", + quality="high", + user="end-user-1", + background="transparent", + moderation="low", + output_compression=50, + ) + + assert optional_params == {"num_images": 1} + def test_flux2_flex_model_info(): model_info = litellm.get_model_info( From 32e6a86319ea9967c4af7a00e26411fd91046871 Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 12 Sep 2026 01:00:44 +0000 Subject: [PATCH 052/102] fix(router): report null cost for unpriced deployments instead of 0 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 5 +- litellm/router.py | 11 ++++- litellm/utils.py | 16 +++++++ .../proxy_server/test_routes_model_info.py | 37 +++++++++++++++ tests/test_litellm/test_router.py | 47 +++++++++++++++++++ 5 files changed, 113 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b5f4236d22..f21c907c292 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -154,7 +154,7 @@ from litellm.types.utils import ( TextCompletionResponse, TokenCountResponse, ) -from litellm.utils import load_credentials_from_list +from litellm.utils import cost_map_omits_token_price, load_credentials_from_list if TYPE_CHECKING: from aiohttp import ClientSession @@ -13618,9 +13618,10 @@ def _enrich_model_info_with_litellm_data( discovered_model_info: Final = ( llm_router.get_discovered_model_info(model_info.get("id")) if llm_router is not None else MappingProxyType({}) ) + unpriced: Final = cost_map_omits_token_price(model_info.get("id"), litellm_model_info.get("key")) for k, v in MappingProxyType({**litellm_model_info, **discovered_model_info}).items(): if k not in model_info or (model_info[k] is None and k in discovered_model_info): - model_info[k] = v + model_info[k] = None if unpriced and k in ("input_cost_per_token", "output_cost_per_token") else v model["model_info"] = model_info # don't return the api key / vertex credentials # don't return the llm credentials diff --git a/litellm/router.py b/litellm/router.py index eec393d6cef..8a543373692 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10575,11 +10575,12 @@ class Router: 2. If not, check if litellm model name is in model info 3. If not, return None """ - from litellm.utils import _update_dictionary + from litellm.utils import _update_dictionary, cost_map_omits_token_price model_info: ModelInfo | None = None custom_model_info: dict | None = None litellm_model_name_model_info: ModelInfo | None = None + base_model_key: str | None = None try: custom_model_info = ( @@ -10606,6 +10607,7 @@ class Router: ## update litellm model info with base model info base_model_info: Final = copy.deepcopy(litellm.get_model_info(model=base_model)) if base_model_info is not None: + base_model_key = base_model_info.get("key") # Base model provides defaults, custom model info overrides custom_model_info = _update_dictionary( cast(dict, base_model_info), @@ -10633,6 +10635,13 @@ class Router: # custom_model_info already includes base_model defaults at this point, if applicable model_info = cast(ModelInfo, custom_model_info) + if model_info is None: + return None + builtin_key: Final = ( + litellm_model_name_model_info.get("key") if litellm_model_name_model_info is not None else None + ) + if cost_map_omits_token_price(model_id, builtin_key, base_model_key): + return cast(ModelInfo, {**model_info, "input_cost_per_token": None, "output_cost_per_token": None}) return model_info def _set_model_group_info(self, model_group: str, user_facing_model_group_name: str) -> ModelGroupInfo | None: diff --git a/litellm/utils.py b/litellm/utils.py index 2c9200fbad7..ec11895a8b5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3156,6 +3156,22 @@ def reapply_runtime_model_cost_registrations() -> None: register_model(model_cost=dict(_runtime_registered_model_cost)) # mutable-ok: snapshot, replay rewrites it +def cost_map_omits_token_price(*keys: object) -> bool: + """Whether the raw ``litellm.model_cost`` entries under ``keys`` exist but none carries a per-token price. + + ``get_model_info`` substitutes 0 for a missing price, which reads exactly like a declared + zero. Surfaces that report pricing use this to keep an unpriced deployment at ``None``. + """ + entries: Final = tuple( + entry + for entry in (litellm.model_cost.get(key) for key in keys if isinstance(key, str)) + if isinstance(entry, dict) + ) + return len(entries) > 0 and not any( + "input_cost_per_token" in entry or "output_cost_per_token" in entry for entry in entries + ) + + def register_model( model_cost: str | dict, *, diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index a1cf838ab6b..203a237113f 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -286,6 +286,43 @@ def test_get_proxy_model_info_surfaces_supports_parallel_function_calling(local_ assert enriched["model_info"]["supports_parallel_function_calling"] is True +def test_model_info_reports_null_cost_for_unpriced_deployment_and_zero_for_declared_zero(): + """A deployment configured with no cost fields must not surface the 0 that ``get_model_info`` + defaults to, since the zero-cost budget bypass only honours a declared zero. The declared zero + and a catalog price still come through.""" + import litellm + + router = litellm.Router( + model_list=[ + { + "model_name": "vllm-unpriced", + "litellm_params": {"model": "openai/vllm-unpriced", "api_key": "x", "api_base": "http://vllm"}, + }, + { + "model_name": "vllm-free", + "litellm_params": { + "model": "openai/vllm-free", + "api_key": "x", + "api_base": "http://vllm", + "input_cost_per_token": 0, + "output_cost_per_token": 0, + }, + }, + {"model_name": "gpt-priced", "litellm_params": {"model": "gpt-4o", "api_key": "x"}}, + ] + ) + + def enriched_cost(model_name: str) -> tuple: + deployment = router.get_model_list(model_name=model_name)[0] + info = proxy_server._enrich_model_info_with_litellm_data({**deployment, "model_info": dict(deployment["model_info"])})["model_info"] + return info.get("input_cost_per_token"), info.get("output_cost_per_token") + + assert enriched_cost("vllm-unpriced") == (None, None) + assert enriched_cost("vllm-free") == (0, 0) + input_cost, output_cost = enriched_cost("gpt-priced") + assert input_cost > 0 and output_cost > 0 + + def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch): from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth from litellm.proxy.auth import model_checks diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 26f9803a022..ceb8fcf6fd1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1888,6 +1888,53 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): assert result.output_cost_per_token is None +def test_model_group_info_cost_none_for_unpriced_deployment_but_zero_when_declared(): + """A deployment with no cost fields anywhere must report None, not the 0 that + get_model_info defaults to, so the reported price matches what the zero-cost + budget bypass accepts. A deployment declaring 0 keeps reporting 0.""" + router = litellm.Router( + model_list=[ + { + "model_name": "vllm-unpriced", + "litellm_params": { + "model": "openai/my-vllm-unpriced", + "api_key": "fake", + "api_base": "http://localhost:8000/v1", + }, + }, + { + "model_name": "vllm-free", + "litellm_params": { + "model": "openai/my-vllm-free", + "api_key": "fake", + "api_base": "http://localhost:8000/v1", + "input_cost_per_token": 0, + "output_cost_per_token": 0, + }, + }, + { + "model_name": "gpt-priced", + "litellm_params": {"model": "gpt-4o", "api_key": "fake"}, + }, + ] + ) + + unpriced = router.get_model_group_info(model_group="vllm-unpriced") + assert unpriced is not None + assert unpriced.input_cost_per_token is None + assert unpriced.output_cost_per_token is None + + free = router.get_model_group_info(model_group="vllm-free") + assert free is not None + assert free.input_cost_per_token == 0 + assert free.output_cost_per_token == 0 + + priced = router.get_model_group_info(model_group="gpt-priced") + assert priced is not None + assert priced.input_cost_per_token is not None and priced.input_cost_per_token > 0 + assert priced.output_cost_per_token is not None and priced.output_cost_per_token > 0 + + @pytest.mark.parametrize( "value,expected", [ From 170581e7101ba821a6bd6a2cb039718f899fc85d Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 12 Sep 2026 01:04:40 +0000 Subject: [PATCH 053/102] fix(router): annotate cast for unpriced model info Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 8a543373692..9ff23c597f1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10641,7 +10641,9 @@ class Router: litellm_model_name_model_info.get("key") if litellm_model_name_model_info is not None else None ) if cost_map_omits_token_price(model_id, builtin_key, base_model_key): - return cast(ModelInfo, {**model_info, "input_cost_per_token": None, "output_cost_per_token": None}) + return cast( # cast-ok: TypedDict spread with overridden keys loses its type + ModelInfo, {**model_info, "input_cost_per_token": None, "output_cost_per_token": None} + ) return model_info def _set_model_group_info(self, model_group: str, user_facing_model_group_name: str) -> ModelGroupInfo | None: From 8549189cbbedc85ed067da8a247512cbac54122f Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 19:59:24 +0000 Subject: [PATCH 054/102] test(proxy): use local_model_cost_map fixture for alias listing test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_utils.py | 37 +++++++++----------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index a79e0798e29..25f4886c6f9 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -2243,32 +2243,27 @@ def test_create_model_info_response_resolves_mode_through_deployment_model(): {"team-embeddings": {"model": "my-embeddings", "hidden": False}}, ], ) -def test_create_model_info_response_resolves_model_group_alias_to_target(model_group_alias): +def test_create_model_info_response_resolves_model_group_alias_to_target(model_group_alias, local_model_cost_map): """A `model_group_alias` row must report the metadata of the group it points at, not the cost-map generalization or nothing that the alias name resolves to.""" from litellm import Router - saved_model_cost = dict(litellm.model_cost) - try: - router = Router( - model_list=[ - { - "model_name": "my-embeddings", - "litellm_params": {"model": "openai/text-embedding-3-small"}, - } - ], - model_group_alias=model_group_alias, - ) + router = Router( + model_list=[ + { + "model_name": "my-embeddings", + "litellm_params": {"model": "openai/text-embedding-3-small"}, + } + ], + model_group_alias=model_group_alias, + ) - alias_response = create_model_info_response( - model_id="team-embeddings", provider="openai", llm_router=router - ) - target_response = create_model_info_response( - model_id="my-embeddings", provider="openai", llm_router=router - ) - finally: - litellm.model_cost.clear() - litellm.model_cost.update(saved_model_cost) + alias_response = create_model_info_response( + model_id="team-embeddings", provider="openai", llm_router=router + ) + target_response = create_model_info_response( + model_id="my-embeddings", provider="openai", llm_router=router + ) assert alias_response["id"] == "team-embeddings" for field in ("mode", "max_input_tokens", "max_output_tokens"): From 657cf18aea50d8825b629d95f9e74467d17391d9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 13:01:55 -0700 Subject: [PATCH 055/102] feat(rust): port exception_type to litellm-core-utils Adds a Rust port of litellm_core_utils/exception_mapping_utils.exception_type with its provider rule tables (OpenAI-compatible, Cohere, Vertex AI), the secret redaction patterns from secret_redaction.py, and a Python repr helper for messages that quote caller values. The public failure shapes are pinned by golden JSON fixtures under tests/test_litellm/rust_bridge/fixtures/public_failures, which the Python side reads too once the bridge is wired to this module. Nothing calls the port yet. --- litellm-rust/Cargo.lock | 29 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/core-utils/Cargo.toml | 5 + .../src/exception_mapping_utils/cohere.rs | 232 ++++++ .../src/exception_mapping_utils/mod.rs | 734 ++++++++++++++++++ .../src/exception_mapping_utils/openai.rs | 586 ++++++++++++++ .../src/exception_mapping_utils/original.rs | 49 ++ .../src/exception_mapping_utils/public.rs | 262 +++++++ .../src/exception_mapping_utils/rules.rs | 403 ++++++++++ .../src/exception_mapping_utils/status.rs | 153 ++++ .../src/exception_mapping_utils/vertex_ai.rs | 574 ++++++++++++++ litellm-rust/crates/core-utils/src/lib.rs | 3 + .../crates/core-utils/src/python_repr.rs | 93 +++ .../crates/core-utils/src/secret_redaction.rs | 97 +++ .../fixtures/public_failures/api.json | 13 + .../public_failures/api_connection.json | 11 + .../status_authentication.json | 13 + .../public_failures/status_bad_gateway.json | 28 + .../public_failures/status_bad_request.json | 28 + .../status_content_policy_violation.json | 28 + .../status_context_window_exceeded.json | 28 + .../status_internal_server.json | 19 + .../public_failures/status_not_found.json | 28 + .../status_permission_denied.json | 19 + .../public_failures/status_rate_limit.json | 28 + .../status_service_unavailable.json | 28 + .../status_unsupported_params.json | 28 + .../public_failures/timeout_with_status.json | 12 + .../timeout_without_status.json | 12 + 29 files changed, 3544 insertions(+) create mode 100644 litellm-rust/crates/core-utils/src/exception_mapping_utils/cohere.rs create mode 100644 litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs create mode 100644 litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs create mode 100644 litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs create mode 100644 litellm-rust/crates/core-utils/src/exception_mapping_utils/public.rs create mode 100644 litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs create mode 100644 litellm-rust/crates/core-utils/src/exception_mapping_utils/status.rs create mode 100644 litellm-rust/crates/core-utils/src/exception_mapping_utils/vertex_ai.rs create mode 100644 litellm-rust/crates/core-utils/src/python_repr.rs create mode 100644 litellm-rust/crates/core-utils/src/secret_redaction.rs create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/api.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/api_connection.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_authentication.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_gateway.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_request.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_content_policy_violation.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_context_window_exceeded.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_internal_server.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_not_found.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_permission_denied.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_rate_limit.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_service_unavailable.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_unsupported_params.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_with_status.json create mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_without_status.json diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 3b9ff62fbaa..6c524124550 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -559,6 +559,21 @@ dependencies = [ "vsimd", ] +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + [[package]] name = "bitflags" version = "2.13.1" @@ -1166,6 +1181,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fancy-regex" +version = "0.19.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d301f5bf187b3c295fce6468d3875037a0bccc5f6b151c63cac2f85babf21912" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fastrand" version = "2.5.0" @@ -2059,11 +2085,14 @@ dependencies = [ name = "litellm-core-utils" version = "0.1.0" dependencies = [ + "fancy-regex", "litellm-types", + "rstest", "serde", "serde_json", "serde_path_to_error", "serde_with", + "strum", "thiserror 2.0.19", "url", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index f8377138050..ffdbf64bb49 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -50,6 +50,7 @@ strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" time = { version = "0.3.53", features = ["parsing"] } criterion = "0.8.2" +fancy-regex = "0.19.2" veil = "0.3.0" [profile.release] diff --git a/litellm-rust/crates/core-utils/Cargo.toml b/litellm-rust/crates/core-utils/Cargo.toml index 109c3312727..59e0ee1a09d 100644 --- a/litellm-rust/crates/core-utils/Cargo.toml +++ b/litellm-rust/crates/core-utils/Cargo.toml @@ -6,10 +6,15 @@ license.workspace = true repository.workspace = true [dependencies] +fancy-regex.workspace = true litellm-types.workspace = true serde.workspace = true serde_json.workspace = true serde_path_to_error = "0.1" serde_with.workspace = true +strum.workspace = true thiserror.workspace = true url.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/cohere.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/cohere.rs new file mode 100644 index 00000000000..381ebd7ad27 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/cohere.rs @@ -0,0 +1,232 @@ +use super::Mapping; +use super::public::{PublicFailure, StatusClass}; +use super::rules::{Kind, ResponseChoice, Rule, apply, contains_any}; + +const fn with_response(class: StatusClass) -> Kind { + Kind::Status { + class, + response: ResponseChoice::Provider, + } +} + +fn original(mapping: &Mapping<'_>) -> String { + format!("CohereException - {}", mapping.original.message) +} + +fn status_is(mapping: &Mapping<'_>, statuses: &[u16]) -> bool { + mapping + .original + .status + .is_some_and(|status| statuses.contains(&status)) +} + +/// `_map_cohere_exception`, in its branch order. A failure no rule claims falls through to +/// the status table. +const RULES: &[Rule] = &[ + Rule { + when: |mapping| { + contains_any( + &mapping.error_str, + &["invalid api token", "No API key provided."], + ) + }, + kind: with_response(StatusClass::Authentication), + message: original, + debug: false, + }, + Rule { + when: |mapping| mapping.error_str.contains("invalid type: parameter"), + kind: with_response(StatusClass::BadRequest), + message: original, + debug: false, + }, + Rule { + when: |mapping| mapping.error_str.contains("too many tokens"), + kind: with_response(StatusClass::ContextWindowExceeded), + message: original, + debug: false, + }, + Rule { + when: |mapping| { + mapping + .error_str + .to_lowercase() + .contains("internal server error") + }, + kind: with_response(StatusClass::InternalServer), + message: |mapping| format!("CohereException - {}", mapping.error_str), + debug: false, + }, + Rule { + when: |mapping| status_is(mapping, &[400, 498]), + kind: with_response(StatusClass::BadRequest), + message: original, + debug: false, + }, + Rule { + when: |mapping| status_is(mapping, &[408]), + kind: Kind::Timeout(None), + message: original, + debug: false, + }, + Rule { + when: |mapping| status_is(mapping, &[500]), + kind: with_response(StatusClass::InternalServer), + message: original, + debug: false, + }, +]; + +pub(super) fn map(mapping: &Mapping<'_>) -> Option { + apply(RULES, mapping).map(|failure| PublicFailure { + llm_provider: Some("cohere".to_string()), + ..failure + }) +} + +#[cfg(test)] +mod tests { + use super::super::testing::{context, failure, http, status, upstream}; + use super::super::{ExceptionFamily, OriginalException, PublicKind}; + use super::*; + + fn mapped(provider: &str, original: &OriginalException) -> Option { + let context = context(provider, ExceptionFamily::Cohere); + map(&Mapping::new(&context, original)) + } + + fn cohere(class: StatusClass, status_code: u16, body: &str, message: &str) -> PublicFailure { + failure( + status(class, upstream(status_code, body)), + message, + "cohere", + ) + } + + #[rstest::rstest] + #[case::invalid_token( + 500, + "invalid api token", + cohere( + StatusClass::Authentication, + 500, + "invalid api token", + "CohereException - invalid api token" + ) + )] + #[case::no_api_key( + 500, + "No API key provided.", + cohere( + StatusClass::Authentication, + 500, + "No API key provided.", + "CohereException - No API key provided." + ) + )] + #[case::invalid_parameter( + 500, + "invalid type: parameter x", + cohere( + StatusClass::BadRequest, + 500, + "invalid type: parameter x", + "CohereException - invalid type: parameter x" + ) + )] + #[case::too_many_tokens( + 500, + "too many tokens", + cohere( + StatusClass::ContextWindowExceeded, + 500, + "too many tokens", + "CohereException - too many tokens" + ) + )] + #[case::internal_server_text( + 400, + "Internal Server Error", + cohere( + StatusClass::InternalServer, + 400, + "Internal Server Error", + "CohereException - Internal Server Error" + ) + )] + #[case::bad_request( + 400, + "rejected", + cohere(StatusClass::BadRequest, 400, "rejected", "CohereException - rejected") + )] + #[case::invalid_token_status( + 498, + "rejected", + cohere(StatusClass::BadRequest, 498, "rejected", "CohereException - rejected") + )] + #[case::request_timeout(408, "rejected", failure(PublicKind::Timeout { status: None }, "CohereException - rejected", "cohere"))] + #[case::internal_server( + 500, + "rejected", + cohere( + StatusClass::InternalServer, + 500, + "rejected", + "CohereException - rejected" + ) + )] + fn each_rule_maps_and_reports_cohere( + #[case] status_code: u16, + #[case] body: &str, + #[case] expected: PublicFailure, + ) { + assert_eq!(mapped("azure_ai", &http(status_code, body)), Some(expected)); + } + + #[rstest::rstest] + #[case::unmapped_status(409)] + #[case::unauthorized(401)] + fn statuses_without_a_rule_fall_through(#[case] status_code: u16) { + assert_eq!(mapped("cohere", &http(status_code, "rejected")), None); + } + + #[test] + fn the_internal_server_rule_uses_the_redacted_text() { + let body = "internal server error Bearer abcdefghijklmnop"; + assert_eq!( + mapped("cohere", &http(400, body)), + Some(cohere( + StatusClass::InternalServer, + 400, + body, + "CohereException - internal server error REDACTED" + )) + ); + } + + #[rstest::rstest] + #[case::token_before_parameter( + "invalid api token invalid type: parameter", + StatusClass::Authentication + )] + #[case::parameter_before_tokens( + "invalid type: parameter too many tokens", + StatusClass::BadRequest + )] + #[case::tokens_before_internal( + "too many tokens Internal Server Error", + StatusClass::ContextWindowExceeded + )] + #[case::internal_before_status("Internal Server Error", StatusClass::InternalServer)] + fn the_earlier_rule_wins_when_two_apply(#[case] body: &str, #[case] class: StatusClass) { + assert_eq!( + mapped("cohere", &http(400, body)), + Some(cohere( + class, + 400, + body, + &format!("CohereException - {body}") + )) + ); + } +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs new file mode 100644 index 00000000000..8892fba1399 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs @@ -0,0 +1,734 @@ +//! A port of Python's `exception_type` for the routes that run in Rust. +//! +//! KNOWN_GAPS: differences from the Python mapper that no Rust route can reach today. Each +//! one stops being acceptable at its trigger. +//! - The Vertex partner-model API base for "claude" models is not built into +//! `extra_information`. Trigger: a Vertex route whose models include Anthropic partner +//! models; then `api_base` gets that branch and a table row. +//! - Python reports the provider `get_llm_provider` resolves for a stripped model name when +//! that name happens to be in the model cost map. Trigger: a route whose model names +//! overlap the cost map; that needs the provider resolution port, not a classifier change. +//! - The generic `APIConnectionError` fallback appends `traceback.format_exc()` to the +//! message. Rust has no Python traceback and does not invent one; a sweep row that reaches +//! it compares the message before the traceback. + +use super::secret_redaction::{redact_string, secret_redaction_enabled}; + +mod cohere; +mod openai; +mod original; +mod public; +mod rules; +mod status; +mod vertex_ai; + +pub use original::{ExceptionFamily, LocalClass, OriginalException}; +pub use public::{HttpStub, PublicFailure, PublicKind, ResponseArg, StatusClass, UpstreamResponse}; + +const DOCS_URL: &str = "https://docs.litellm.ai/docs"; + +const TIMEOUT_MARKERS: &[&str] = &[ + "Request Timeout Error", + "Request timed out", + "Timed out generating response", + "The read operation timed out", +]; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ExceptionContext { + pub model: String, + pub custom_llm_provider: Option, + pub family: ExceptionFamily, + pub asynchronous: bool, + pub suppress_debug_info: bool, + pub redact_messages_in_exceptions: bool, + pub vertex_project: Option, + pub vertex_location: Option, + pub model_group: Option, + pub deployment: Option, + pub user_api_key_alias: Option, + pub user_api_key_team_alias: Option, +} + +/// The attributes `exception_type` reads off the Python exception: a provider error +/// (`BaseLLMException`) carries a status, a response and a request, a plain exception +/// carries only its text. +struct Raised { + status: Option, + status_is_synthesized: bool, + message: String, + response: Option, +} + +impl Raised { + fn provider( + status: u16, + message: String, + body: String, + headers: Vec<(String, String)>, + ) -> Self { + Self { + status: Some(status), + status_is_synthesized: false, + message, + response: Some(UpstreamResponse { + status, + body, + headers, + }), + } + } + + fn plain(message: String) -> Self { + Self { + status: None, + status_is_synthesized: false, + message, + response: None, + } + } + + fn new(original: &OriginalException, asynchronous: bool) -> Self { + match original { + OriginalException::Http { + status, + body, + headers, + } => Self::provider(*status, body.clone(), body.clone(), headers.clone()), + OriginalException::Connection { message } => Self { + status_is_synthesized: true, + ..Self::provider(500, message.clone(), String::new(), Vec::new()) + }, + OriginalException::Timeout { + timeout_seconds, + elapsed_seconds, + } => Self::provider( + 408, + timeout_message(asynchronous, *timeout_seconds, *elapsed_seconds), + String::new(), + Vec::new(), + ), + OriginalException::Response { message } + | OriginalException::Local { message, .. } + | OriginalException::Public { message, .. } => Self::plain(message.clone()), + } + } +} + +/// The text `litellm.Timeout` carries when the Python HTTP handler times out: the sync +/// and async handlers word it differently. +fn timeout_message( + asynchronous: bool, + timeout_seconds: Option, + elapsed_seconds: Option, +) -> String { + let timeout = python_float(timeout_seconds); + if asynchronous { + let elapsed = + python_float(elapsed_seconds.map(|seconds| (seconds * 1000.0).round() / 1000.0)); + format!( + "litellm.Timeout: Connection timed out. Timeout passed={timeout}, time taken={elapsed} seconds" + ) + } else { + format!("litellm.Timeout: Connection timed out after {timeout} seconds.") + } +} + +fn python_float(value: Option) -> String { + match value { + None => "None".to_string(), + Some(value) if value.fract() == 0.0 => format!("{value:.1}"), + Some(value) => value.to_string(), + } +} + +/// Everything the rules read: the original as Python sees it and the text `exception_type` +/// derives from the context before any provider mapper runs. +struct Mapping<'a> { + context: &'a ExceptionContext, + original: Raised, + provider: &'a str, + error_str: String, + exception_provider: String, + extra_information: String, +} + +impl<'a> Mapping<'a> { + fn new(context: &'a ExceptionContext, original: &OriginalException) -> Self { + let original = Raised::new(original, context.asynchronous); + let error_str = if secret_redaction_enabled() { + redact_string(&original.message) + } else { + original.message.clone() + }; + Self { + context, + original, + provider: context.custom_llm_provider.as_deref().unwrap_or_default(), + error_str, + exception_provider: match &context.custom_llm_provider { + None => "None".to_string(), + Some(provider) => exception_provider(provider), + }, + extra_information: extra_information(context, api_base(context).as_deref()), + } + } + + fn failure(&self, kind: PublicKind, message: String, debug: bool) -> PublicFailure { + PublicFailure { + kind, + message, + model: self.context.model.clone(), + llm_provider: self.context.custom_llm_provider.clone(), + litellm_debug_info: debug.then(|| self.extra_information.clone()), + litellm_response_headers: None, + print_banner: false, + } + } +} + +pub fn exception_type(context: &ExceptionContext, original: &OriginalException) -> PublicFailure { + if let OriginalException::Public { class, message } = original { + return PublicFailure { + kind: PublicKind::Status { + status_class: *class, + response: None, + }, + message: message.clone(), + model: context.model.clone(), + llm_provider: context.custom_llm_provider.clone(), + litellm_debug_info: None, + litellm_response_headers: None, + print_banner: false, + }; + } + let mapping = Mapping::new(context, original); + let litellm_response_headers = mapping + .original + .response + .as_ref() + .map(|response| response.headers.clone()) + .filter(|headers| !headers.is_empty()); + PublicFailure { + litellm_response_headers, + print_banner: !context.suppress_debug_info, + ..map(&mapping) + } +} + +fn map(mapping: &Mapping<'_>) -> PublicFailure { + if rules::contains_any(&mapping.error_str, TIMEOUT_MARKERS) { + return mapping.failure( + PublicKind::Timeout { status: None }, + format!( + "APITimeoutError - Request timed out. Error_str: {}", + mapping.error_str + ), + true, + ); + } + let provider_failure = match mapping.context.family { + ExceptionFamily::OpenAiCompatible => openai::map(mapping), + ExceptionFamily::VertexAi => vertex_ai::map(mapping), + ExceptionFamily::Cohere => cohere::map(mapping), + ExceptionFamily::Other => None, + }; + provider_failure + .or_else(|| status::map(mapping)) + .unwrap_or_else(|| unmapped(mapping)) +} + +/// The `APIConnectionError` Python raises when no mapper claimed the failure: with the +/// provider prefix for a provider error, with the bare text for a plain exception. +fn unmapped(mapping: &Mapping<'_>) -> PublicFailure { + let message = match mapping.original.status { + Some(_) => format!("{} - {}", mapping.exception_provider, mapping.error_str), + None => mapping.original.message.clone(), + }; + mapping.failure(PublicKind::ApiConnection, message, false) +} + +fn exception_provider(provider: &str) -> String { + let mut characters = provider.chars(); + match characters.next() { + Some(first) => format!("{}{}Exception", first.to_uppercase(), characters.as_str()), + None => String::new(), + } +} + +fn python_capitalize(value: &str) -> String { + let mut characters = value.chars(); + match characters.next() { + Some(first) => format!( + "{}{}", + first.to_uppercase(), + characters.as_str().to_lowercase() + ), + None => String::new(), + } +} + +fn api_base(context: &ExceptionContext) -> Option { + match (&context.vertex_location, &context.vertex_project) { + (Some(location), Some(project)) => Some(format!( + "{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/google/models/{}:generateContent", + context.model + )), + _ => None, + } +} + +fn extra_information(context: &ExceptionContext, api_base: Option<&str>) -> String { + let lines = [ + Some(format!("\nModel: {}", context.model)), + api_base.map(|api_base| format!("\nAPI Base: `{api_base}`")), + (!context.redact_messages_in_exceptions).then(|| "\nMessages: `None`".to_string()), + context + .model_group + .as_ref() + .map(|value| format!("\nmodel_group: `{value}`\n")), + context + .deployment + .as_ref() + .map(|value| format!("\ndeployment: `{value}`\n")), + context + .vertex_project + .as_ref() + .map(|value| format!("\nvertex_project: `{value}`\n")), + context + .vertex_location + .as_ref() + .map(|value| format!("\nvertex_location: `{value}`\n")), + ]; + let information: String = lines.into_iter().flatten().collect(); + match &context.user_api_key_alias { + Some(alias) => format!( + "\n\nKey Name: `{alias}`\nTeam: `{}`{information}", + context.user_api_key_team_alias.as_deref().unwrap_or("None") + ), + None => information, + } +} + +#[cfg(test)] +mod testing { + use super::*; + + pub(super) const DEBUG: &str = "\nModel: ocr-model\nMessages: `None`"; + + pub(super) fn context(provider: &str, family: ExceptionFamily) -> ExceptionContext { + ExceptionContext { + model: "ocr-model".into(), + custom_llm_provider: Some(provider.into()), + family, + suppress_debug_info: true, + ..ExceptionContext::default() + } + } + + pub(super) fn http(status: u16, body: &str) -> OriginalException { + OriginalException::Http { + status, + body: body.into(), + headers: vec![("retry-after".into(), "7".into())], + } + } + + pub(super) fn upstream(status: u16, body: &str) -> Option { + Some(ResponseArg::Upstream(UpstreamResponse { + status, + body: body.into(), + headers: vec![("retry-after".into(), "7".into())], + })) + } + + pub(super) fn status(class: StatusClass, response: Option) -> PublicKind { + PublicKind::Status { + status_class: class, + response, + } + } + + /// The failure a rule builds before `exception_type` adds the response headers and the + /// banner flag. + pub(super) fn failure(kind: PublicKind, message: &str, provider: &str) -> PublicFailure { + PublicFailure { + kind, + message: message.into(), + model: "ocr-model".into(), + llm_provider: Some(provider.into()), + litellm_debug_info: None, + litellm_response_headers: None, + print_banner: false, + } + } + + pub(super) fn with_debug(failure: PublicFailure) -> PublicFailure { + PublicFailure { + litellm_debug_info: Some(DEBUG.into()), + ..failure + } + } + + /// What `exception_type` returns for an `http` original the rule mapped to `failure`. + pub(super) fn with_headers(failure: PublicFailure) -> PublicFailure { + PublicFailure { + litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), + ..failure + } + } +} + +#[cfg(test)] +mod tests { + use super::testing::{DEBUG, context, failure, http, status, upstream, with_debug}; + use super::*; + + fn openai() -> ExceptionContext { + context("mistral", ExceptionFamily::OpenAiCompatible) + } + + #[test] + fn a_public_original_passes_through_without_banner_debug_or_prefix() { + let original = OriginalException::Public { + class: StatusClass::UnsupportedParams, + message: "Invalid `req_format`".into(), + }; + let context = ExceptionContext { + suppress_debug_info: false, + ..openai() + }; + assert_eq!( + exception_type(&context, &original), + failure( + status(StatusClass::UnsupportedParams, None), + "Invalid `req_format`", + "mistral" + ) + ); + } + + #[rstest::rstest] + #[case::vertex_family_status_rule(ExceptionFamily::VertexAi, "vertex_ai", PublicFailure { + litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), + ..with_debug(failure(status(StatusClass::BadRequest, upstream(409, "rejected")), "Vertex_aiException - rejected", "vertex_ai")) + })] + #[case::cohere_family(ExceptionFamily::Cohere, "cohere", PublicFailure { + litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), + ..with_debug(failure(status(StatusClass::BadRequest, upstream(409, "rejected")), "CohereException - rejected", "cohere")) + })] + #[case::other_family(ExceptionFamily::Other, "reducto", PublicFailure { + litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), + ..with_debug(failure(status(StatusClass::BadRequest, upstream(409, "rejected")), "ReductoException - rejected", "reducto")) + })] + fn families_without_a_409_rule_reach_the_status_table( + #[case] family: ExceptionFamily, + #[case] provider: &str, + #[case] expected: PublicFailure, + ) { + assert_eq!( + exception_type(&context(provider, family), &http(409, "rejected")), + expected + ); + } + + #[test] + fn the_openai_family_claims_a_409_before_the_status_table() { + assert_eq!( + exception_type(&openai(), &http(409, "rejected")), + PublicFailure { + litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), + ..with_debug(failure( + PublicKind::Api { + status: 409, + request_url: DOCS_URL + }, + "APIError: MistralException - rejected", + "mistral" + )) + } + ); + } + + #[rstest::rstest] + #[case::request_timeout_error("Request Timeout Error")] + #[case::request_timed_out("Request timed out")] + #[case::timed_out_generating("Timed out generating response")] + #[case::read_operation("The read operation timed out")] + fn timeout_markers_win_over_every_family(#[case] marker: &str) { + let body = format!("rate limit {marker}"); + for family in [ + ExceptionFamily::OpenAiCompatible, + ExceptionFamily::VertexAi, + ExceptionFamily::Cohere, + ExceptionFamily::Other, + ] { + assert_eq!( + exception_type(&context("mistral", family), &http(429, &body)), + PublicFailure { + litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), + ..with_debug(failure( + PublicKind::Timeout { status: None }, + &format!("APITimeoutError - Request timed out. Error_str: {body}"), + "mistral" + )) + } + ); + } + } + + #[rstest::rstest] + #[case::provider_error_keeps_the_prefix(http(409, "rejected"), "ReductoException - rejected")] + #[case::synthesized_status_skips_the_status_table( + OriginalException::Connection { message: "refused".into() }, + "ReductoException - refused" + )] + #[case::plain_exception_keeps_its_text( + OriginalException::Local { class: LocalClass::FileNotFound, message: "File not found: /a".into() }, + "File not found: /a" + )] + fn unmapped_failures_are_connection_errors( + #[case] original: OriginalException, + #[case] message: &str, + ) { + let context = context("reducto", ExceptionFamily::Other); + let expected = failure(PublicKind::ApiConnection, message, "reducto"); + let actual = exception_type(&context, &original); + assert_eq!( + PublicFailure { + litellm_response_headers: None, + ..actual + }, + expected + ); + } + + #[test] + fn a_missing_provider_renders_like_python_none() { + let context = ExceptionContext { + custom_llm_provider: None, + family: ExceptionFamily::Other, + ..openai() + }; + assert_eq!( + exception_type(&context, &http(401, "rejected")), + PublicFailure { + llm_provider: None, + litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), + ..with_debug(failure( + status(StatusClass::Authentication, upstream(401, "rejected")), + "None - rejected", + "unused" + )) + } + ); + } + + #[rstest::rstest] + #[case::suppressed(true, false)] + #[case::printed(false, true)] + fn the_banner_prints_unless_debug_info_is_suppressed( + #[case] suppress_debug_info: bool, + #[case] print_banner: bool, + ) { + let context = ExceptionContext { + suppress_debug_info, + ..openai() + }; + assert_eq!( + exception_type(&context, &http(400, "rejected")).print_banner, + print_banner + ); + } + + #[test] + fn empty_upstream_headers_are_not_reported() { + let original = OriginalException::Http { + status: 400, + body: "rejected".into(), + headers: Vec::new(), + }; + assert_eq!( + exception_type(&openai(), &original).litellm_response_headers, + None + ); + } + + #[test] + fn messages_are_redacted_before_markers_and_prefixes() { + let body = "rejected Bearer abcdefghijklmnop"; + assert_eq!( + exception_type( + &context("reducto", ExceptionFamily::Other), + &http(400, body) + ) + .message, + "ReductoException - rejected REDACTED" + ); + } + + const SYNC_TIMEOUT: &str = "litellm.Timeout: Connection timed out after 0.5 seconds."; + + #[rstest::rstest] + #[case::sync(false, Some(0.5), Some(0.5031), SYNC_TIMEOUT)] + #[case::async_rounds_the_elapsed_time( + true, + Some(0.5), + Some(0.5031), + "litellm.Timeout: Connection timed out. Timeout passed=0.5, time taken=0.503 seconds" + )] + #[case::whole_seconds_keep_a_decimal( + true, + Some(600.0), + Some(2.0), + "litellm.Timeout: Connection timed out. Timeout passed=600.0, time taken=2.0 seconds" + )] + #[case::unknown_values_render_as_none( + true, + None, + None, + "litellm.Timeout: Connection timed out. Timeout passed=None, time taken=None seconds" + )] + fn timeout_text_follows_the_delivery_mode( + #[case] asynchronous: bool, + #[case] timeout_seconds: Option, + #[case] elapsed_seconds: Option, + #[case] expected: &str, + ) { + assert_eq!( + timeout_message(asynchronous, timeout_seconds, elapsed_seconds), + expected + ); + } + + #[rstest::rstest] + #[case::sync(false, SYNC_TIMEOUT)] + #[case::async_( + true, + "litellm.Timeout: Connection timed out. Timeout passed=0.5, time taken=0.503 seconds" + )] + fn a_timeout_is_a_408_carrying_the_handler_text( + #[case] asynchronous: bool, + #[case] text: &str, + ) { + let context = ExceptionContext { + asynchronous, + ..openai() + }; + let original = OriginalException::Timeout { + timeout_seconds: Some(0.5), + elapsed_seconds: Some(0.5031), + }; + assert_eq!( + exception_type(&context, &original), + with_debug(failure( + PublicKind::Timeout { status: None }, + &format!("Timeout Error: MistralException - {text}"), + "mistral" + )) + ); + } + + #[test] + fn a_refused_connection_is_a_500_with_an_empty_response() { + assert_eq!( + exception_type( + &openai(), + &OriginalException::Connection { + message: "refused".into() + } + ), + with_debug(failure( + status( + StatusClass::InternalServer, + Some(ResponseArg::Upstream(UpstreamResponse { + status: 500, + body: String::new(), + headers: Vec::new(), + })) + ), + "InternalServerError: MistralException - refused", + "mistral" + )) + ); + } + + #[test] + fn debug_information_follows_the_python_layout() { + let context = ExceptionContext { + vertex_project: Some("project".into()), + vertex_location: Some("region".into()), + model_group: Some("ocr".into()), + deployment: Some("deployment".into()), + user_api_key_alias: Some("key".into()), + ..openai() + }; + assert_eq!( + extra_information(&context, api_base(&context).as_deref()), + concat!( + "\n\nKey Name: `key`\nTeam: `None`", + "\nModel: ocr-model", + "\nAPI Base: `region-aiplatform.googleapis.com/v1/projects/project/locations/region/publishers/google/models/ocr-model:generateContent`", + "\nMessages: `None`", + "\nmodel_group: `ocr`\n", + "\ndeployment: `deployment`\n", + "\nvertex_project: `project`\n", + "\nvertex_location: `region`\n", + ) + ); + } + + #[rstest::rstest] + #[case::bare(ExceptionContext::default(), "\nModel: ")] + #[case::redacted_messages(ExceptionContext { redact_messages_in_exceptions: true, model: "m".into(), ..ExceptionContext::default() }, "\nModel: m")] + #[case::messages(ExceptionContext { model: "m".into(), ..ExceptionContext::default() }, "\nModel: m\nMessages: `None`")] + #[case::team_alias( + ExceptionContext { model: "m".into(), user_api_key_alias: Some("key".into()), user_api_key_team_alias: Some("team".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, + "\n\nKey Name: `key`\nTeam: `team`\nModel: m" + )] + #[case::team_alias_without_key_is_ignored( + ExceptionContext { model: "m".into(), user_api_key_team_alias: Some("team".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, + "\nModel: m" + )] + #[case::project_without_location_has_no_api_base( + ExceptionContext { model: "m".into(), vertex_project: Some("p".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, + "\nModel: m\nvertex_project: `p`\n" + )] + #[case::location_without_project_has_no_api_base( + ExceptionContext { model: "m".into(), vertex_location: Some("l".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, + "\nModel: m\nvertex_location: `l`\n" + )] + fn each_optional_context_field_adds_its_own_line( + #[case] context: ExceptionContext, + #[case] expected: &str, + ) { + assert_eq!( + extra_information(&context, api_base(&context).as_deref()), + expected + ); + } + + #[rstest::rstest] + #[case::lowercase("mistral", "MistralException")] + #[case::keeps_the_rest("azure_ai", "Azure_aiException")] + #[case::empty("", "")] + fn exception_provider_capitalizes_only_the_first_letter( + #[case] provider: &str, + #[case] expected: &str, + ) { + assert_eq!(exception_provider(provider), expected); + } + + #[rstest::rstest] + #[case::lowers_the_rest("vERTEX_AI", "Vertex_ai")] + #[case::empty("", "")] + fn python_capitalize_lowers_the_rest(#[case] value: &str, #[case] expected: &str) { + assert_eq!(python_capitalize(value), expected); + } + + #[test] + fn debug_constant_matches_the_default_test_context() { + let context = openai(); + assert_eq!(extra_information(&context, None), DEBUG); + } +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs new file mode 100644 index 00000000000..ffbb172582b --- /dev/null +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs @@ -0,0 +1,586 @@ +use super::public::{PublicFailure, StatusClass}; +use super::rules::{ + ApiStatus, Kind, ResponseChoice, Rule, apply, contains_any, is_context_window_exceeded, + is_rate_limit, +}; +use super::{DOCS_URL, Mapping}; + +const OPENAI_URL: &str = "https://api.openai.com/v1"; + +const ENCRYPTED_CONTENT_HELP: &str = "\n\n This error occurs when load balancing Responses API across deployments with different API keys.\n Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n router_settings:\n enable_pre_call_checks: true\n optional_pre_call_checks:\n - encrypted_content_affinity\n\n Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing"; + +const fn with_response(class: StatusClass) -> Kind { + Kind::Status { + class, + response: ResponseChoice::Provider, + } +} + +fn exception_provider(mapping: &Mapping<'_>) -> String { + if mapping.provider == "openai" { + "OpenAIException".to_string() + } else { + super::exception_provider(mapping.provider) + } +} + +/// The raw message with OpenAI's own names swapped for the provider's. +fn message(mapping: &Mapping<'_>) -> String { + let provider = mapping.provider; + mapping + .original + .message + .replace("OPENAI", &provider.to_uppercase()) + .replace("openai.OpenAIError", &format!("{provider}.{provider}Error")) +} + +fn prefixed(mapping: &Mapping<'_>, label: &str) -> String { + format!( + "{label}{} - {}", + exception_provider(mapping), + message(mapping) + ) +} + +fn status_is(mapping: &Mapping<'_>, statuses: &[u16]) -> bool { + mapping + .original + .status + .is_some_and(|status| statuses.contains(&status)) +} + +/// `_map_openai_exception`, in its branch order. +const RULES: &[Rule] = &[ + Rule { + when: |mapping| is_rate_limit(&mapping.error_str, mapping.original.status), + kind: with_response(StatusClass::RateLimit), + message: |mapping| prefixed(mapping, "RateLimitError: "), + debug: false, + }, + Rule { + when: |mapping| is_context_window_exceeded(&mapping.error_str), + kind: with_response(StatusClass::ContextWindowExceeded), + message: |mapping| prefixed(mapping, "ContextWindowExceededError: "), + debug: true, + }, + Rule { + when: |mapping| { + mapping.error_str.contains("invalid_request_error") + && mapping.error_str.contains("model_not_found") + }, + kind: with_response(StatusClass::NotFound), + message: |mapping| prefixed(mapping, ""), + debug: true, + }, + Rule { + when: |mapping| mapping.error_str.contains("A timeout occurred"), + kind: Kind::Timeout(None), + message: |mapping| prefixed(mapping, ""), + debug: true, + }, + Rule { + when: |mapping| { + let error_str = &mapping.error_str; + (error_str.contains("invalid_request_error") + && error_str.contains("content_policy_violation")) + || (error_str.contains("Invalid prompt") + && error_str.contains("violating our usage policy")) + || error_str + .to_lowercase() + .contains("request was rejected as a result of the safety system") + }, + kind: with_response(StatusClass::ContentPolicyViolation), + message: |mapping| prefixed(mapping, "ContentPolicyViolationError: "), + debug: true, + }, + Rule { + when: |mapping| { + contains_any( + &mapping.error_str, + &["invalid_encrypted_content", "could not be verified"], + ) + }, + kind: with_response(StatusClass::BadRequest), + message: |mapping| { + format!( + "{} - {}{ENCRYPTED_CONTENT_HELP}", + exception_provider(mapping), + message(mapping) + ) + }, + debug: true, + }, + Rule { + when: |mapping| { + mapping.error_str.contains("invalid_request_error") + && !mapping.error_str.contains("Incorrect API key provided") + }, + kind: with_response(StatusClass::BadRequest), + message: |mapping| prefixed(mapping, ""), + debug: true, + }, + Rule { + when: |mapping| { + contains_any( + &mapping.error_str, + &[ + "Web server is returning an unknown error", + "The server had an error processing your request.", + ], + ) + }, + kind: Kind::Status { + class: StatusClass::InternalServer, + response: ResponseChoice::Omitted, + }, + message: |mapping| prefixed(mapping, ""), + debug: false, + }, + Rule { + when: |mapping| mapping.error_str.contains("Request too large"), + kind: with_response(StatusClass::RateLimit), + message: |mapping| prefixed(mapping, "RateLimitError: "), + debug: true, + }, + Rule { + when: |mapping| { + mapping.error_str.contains("The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable") + }, + kind: with_response(StatusClass::Authentication), + message: |mapping| prefixed(mapping, "AuthenticationError: "), + debug: true, + }, + Rule { + when: |mapping| { + mapping + .error_str + .contains("Mistral API raised a streaming error") + }, + kind: Kind::Api { + status: ApiStatus::Fixed(500), + request_url: OPENAI_URL, + }, + message: |mapping| prefixed(mapping, ""), + debug: true, + }, + Rule { + when: |mapping| mapping.original.status.is_none(), + kind: Kind::ApiConnection, + message: |mapping| prefixed(mapping, "APIConnectionError: "), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[400, 422]), + kind: with_response(StatusClass::BadRequest), + message: |mapping| prefixed(mapping, ""), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[401]), + kind: with_response(StatusClass::Authentication), + message: |mapping| prefixed(mapping, "AuthenticationError: "), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[404]), + kind: with_response(StatusClass::NotFound), + message: |mapping| prefixed(mapping, "NotFoundError: "), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[408]), + kind: Kind::Timeout(None), + message: |mapping| prefixed(mapping, "Timeout Error: "), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[429]), + kind: with_response(StatusClass::RateLimit), + message: |mapping| prefixed(mapping, "RateLimitError: "), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[500]), + kind: with_response(StatusClass::InternalServer), + message: |mapping| prefixed(mapping, "InternalServerError: "), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[502]), + kind: with_response(StatusClass::BadGateway), + message: |mapping| prefixed(mapping, "BadGatewayError: "), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[503]), + kind: with_response(StatusClass::ServiceUnavailable), + message: |mapping| prefixed(mapping, "ServiceUnavailableError: "), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, &[504]), + kind: Kind::Timeout(Some(504)), + message: |mapping| prefixed(mapping, "Timeout Error: "), + debug: true, + }, + Rule { + when: |_| true, + kind: Kind::Api { + status: ApiStatus::Original, + request_url: DOCS_URL, + }, + message: |mapping| prefixed(mapping, "APIError: "), + debug: true, + }, +]; + +pub(super) fn map(mapping: &Mapping<'_>) -> Option { + apply(RULES, mapping) +} + +#[cfg(test)] +mod tests { + use super::super::testing::{context, failure, http, status, upstream, with_debug}; + use super::super::{ExceptionFamily, OriginalException, PublicKind}; + use super::*; + + fn mapped(provider: &str, original: &OriginalException) -> PublicFailure { + let context = context(provider, ExceptionFamily::OpenAiCompatible); + map(&Mapping::new(&context, original)).expect("the OpenAI table ends in a catch-all") + } + + fn kind(class: StatusClass, status_code: u16, body: &str) -> PublicKind { + status(class, upstream(status_code, body)) + } + + #[rstest::rstest] + #[case::rate_limit_phrase( + 400, + "rate limit reached", + failure( + kind(StatusClass::RateLimit, 400, "rate limit reached"), + "RateLimitError: MistralException - rate limit reached", + "mistral", + ) + )] + #[case::context_window( + 500, + "This model's maximum context length is 10", + with_debug(failure( + kind( + StatusClass::ContextWindowExceeded, + 500, + "This model's maximum context length is 10" + ), + "ContextWindowExceededError: MistralException - This model's maximum context length is 10", + "mistral", + )) + )] + #[case::model_not_found( + 400, + "invalid_request_error model_not_found", + with_debug(failure( + kind(StatusClass::NotFound, 400, "invalid_request_error model_not_found"), + "MistralException - invalid_request_error model_not_found", + "mistral", + )) + )] + #[case::timeout_occurred(400, "A timeout occurred", with_debug(failure( + PublicKind::Timeout { status: None }, + "MistralException - A timeout occurred", + "mistral", + )))] + #[case::content_policy_error_code( + 400, + "invalid_request_error content_policy_violation", + with_debug(failure( + kind( + StatusClass::ContentPolicyViolation, + 400, + "invalid_request_error content_policy_violation" + ), + "ContentPolicyViolationError: MistralException - invalid_request_error content_policy_violation", + "mistral", + )) + )] + #[case::content_policy_usage_policy( + 400, + "Invalid prompt violating our usage policy", + with_debug(failure( + kind( + StatusClass::ContentPolicyViolation, + 400, + "Invalid prompt violating our usage policy" + ), + "ContentPolicyViolationError: MistralException - Invalid prompt violating our usage policy", + "mistral", + )) + )] + #[case::content_policy_safety_system( + 400, + "Request was rejected as a result of the safety system", + with_debug(failure( + kind( + StatusClass::ContentPolicyViolation, + 400, + "Request was rejected as a result of the safety system" + ), + "ContentPolicyViolationError: MistralException - Request was rejected as a result of the safety system", + "mistral", + )) + )] + #[case::encrypted_content(400, "invalid_encrypted_content", with_debug(failure( + kind(StatusClass::BadRequest, 400, "invalid_encrypted_content"), + &format!("MistralException - invalid_encrypted_content{ENCRYPTED_CONTENT_HELP}"), + "mistral", + )))] + #[case::unverifiable_content(400, "could not be verified", with_debug(failure( + kind(StatusClass::BadRequest, 400, "could not be verified"), + &format!("MistralException - could not be verified{ENCRYPTED_CONTENT_HELP}"), + "mistral", + )))] + #[case::invalid_request( + 429, + "invalid_request_error bad field", + with_debug(failure( + kind(StatusClass::BadRequest, 429, "invalid_request_error bad field"), + "MistralException - invalid_request_error bad field", + "mistral", + )) + )] + #[case::unknown_server_error( + 400, + "Web server is returning an unknown error", + failure( + status(StatusClass::InternalServer, None), + "MistralException - Web server is returning an unknown error", + "mistral", + ) + )] + #[case::server_had_an_error( + 400, + "The server had an error processing your request.", + failure( + status(StatusClass::InternalServer, None), + "MistralException - The server had an error processing your request.", + "mistral", + ) + )] + #[case::request_too_large( + 400, + "Request too large", + with_debug(failure( + kind(StatusClass::RateLimit, 400, "Request too large"), + "RateLimitError: MistralException - Request too large", + "mistral", + )) + )] + #[case::missing_client_api_key( + 400, + "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable", + with_debug(failure( + kind( + StatusClass::Authentication, + 400, + "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable", + ), + "AuthenticationError: MistralException - The api_key client option must be set either by passing api_key to the client or by setting the MISTRAL_API_KEY environment variable", + "mistral", + )) + )] + #[case::mistral_streaming_error(400, "Mistral API raised a streaming error", with_debug(failure( + PublicKind::Api { status: 500, request_url: OPENAI_URL }, + "MistralException - Mistral API raised a streaming error", + "mistral", + )))] + fn each_text_rule_maps_by_the_body( + #[case] status_code: u16, + #[case] body: &str, + #[case] expected: PublicFailure, + ) { + assert_eq!(mapped("mistral", &http(status_code, body)), expected); + } + + #[rstest::rstest] + #[case::bad_request( + 400, + kind(StatusClass::BadRequest, 400, "rejected"), + "MistralException - rejected" + )] + #[case::unprocessable( + 422, + kind(StatusClass::BadRequest, 422, "rejected"), + "MistralException - rejected" + )] + #[case::authentication( + 401, + kind(StatusClass::Authentication, 401, "rejected"), + "AuthenticationError: MistralException - rejected" + )] + #[case::not_found( + 404, + kind(StatusClass::NotFound, 404, "rejected"), + "NotFoundError: MistralException - rejected" + )] + #[case::request_timeout(408, PublicKind::Timeout { status: None }, "Timeout Error: MistralException - rejected")] + #[case::rate_limited( + 429, + kind(StatusClass::RateLimit, 429, "rejected"), + "RateLimitError: MistralException - rejected" + )] + #[case::internal_server( + 500, + kind(StatusClass::InternalServer, 500, "rejected"), + "InternalServerError: MistralException - rejected" + )] + #[case::bad_gateway( + 502, + kind(StatusClass::BadGateway, 502, "rejected"), + "BadGatewayError: MistralException - rejected" + )] + #[case::service_unavailable( + 503, + kind(StatusClass::ServiceUnavailable, 503, "rejected"), + "ServiceUnavailableError: MistralException - rejected" + )] + #[case::gateway_timeout(504, PublicKind::Timeout { status: Some(504) }, "Timeout Error: MistralException - rejected")] + #[case::any_other_status(409, PublicKind::Api { status: 409, request_url: DOCS_URL }, "APIError: MistralException - rejected")] + fn each_status_rule_maps_by_the_status( + #[case] status_code: u16, + #[case] kind: PublicKind, + #[case] message: &str, + ) { + assert_eq!( + mapped("mistral", &http(status_code, "rejected")), + with_debug(failure(kind, message, "mistral")) + ); + } + + #[test] + fn a_failure_without_a_status_is_a_connection_error() { + let original = OriginalException::Response { + message: "invalid OCR response field: pages".into(), + }; + assert_eq!( + mapped("mistral", &original), + with_debug(failure( + PublicKind::ApiConnection, + "APIConnectionError: MistralException - invalid OCR response field: pages", + "mistral" + )) + ); + } + + #[rstest::rstest] + #[case::rate_limit_before_context_window( + 400, + "rate limit and This model's maximum context length is 10", + kind( + StatusClass::RateLimit, + 400, + "rate limit and This model's maximum context length is 10" + ), + "RateLimitError: MistralException - rate limit and This model's maximum context length is 10", + false + )] + #[case::context_window_before_content_policy( + 400, + "This model's maximum context length is 10 invalid_request_error content_policy_violation", + kind( + StatusClass::ContextWindowExceeded, + 400, + "This model's maximum context length is 10 invalid_request_error content_policy_violation" + ), + "ContextWindowExceededError: MistralException - This model's maximum context length is 10 invalid_request_error content_policy_violation", + true + )] + #[case::model_not_found_before_invalid_request( + 400, + "invalid_request_error model_not_found", + kind(StatusClass::NotFound, 400, "invalid_request_error model_not_found"), + "MistralException - invalid_request_error model_not_found", + true + )] + #[case::timeout_before_invalid_request( + 400, + "A timeout occurred invalid_request_error", + PublicKind::Timeout { status: None }, + "MistralException - A timeout occurred invalid_request_error", + true + )] + #[case::content_policy_before_invalid_request( + 400, + "invalid_request_error content_policy_violation", + kind( + StatusClass::ContentPolicyViolation, + 400, + "invalid_request_error content_policy_violation" + ), + "ContentPolicyViolationError: MistralException - invalid_request_error content_policy_violation", + true + )] + #[case::invalid_request_with_a_bad_key_falls_to_the_status( + 401, + "invalid_request_error Incorrect API key provided", + kind( + StatusClass::Authentication, + 401, + "invalid_request_error Incorrect API key provided" + ), + "AuthenticationError: MistralException - invalid_request_error Incorrect API key provided", + true + )] + #[case::text_rules_before_status( + 429, + "invalid_request_error bad field", + kind(StatusClass::BadRequest, 429, "invalid_request_error bad field"), + "MistralException - invalid_request_error bad field", + true + )] + #[case::echoed_429_is_not_a_rate_limit( + 400, + "token 429 in the prompt", + kind(StatusClass::BadRequest, 400, "token 429 in the prompt"), + "MistralException - token 429 in the prompt", + true + )] + fn the_earlier_rule_wins_when_two_apply( + #[case] status_code: u16, + #[case] body: &str, + #[case] kind: PublicKind, + #[case] message: &str, + #[case] debug: bool, + ) { + let expected = failure(kind, message, "mistral"); + assert_eq!( + mapped("mistral", &http(status_code, body)), + if debug { + with_debug(expected) + } else { + expected + } + ); + } + + #[rstest::rstest] + #[case::provider_names_replace_openai( + "azure_ai", + "OPENAI said openai.OpenAIError", + "Azure_aiException - AZURE_AI said azure_ai.azure_aiError" + )] + #[case::openai_keeps_its_own_name("openai", "rejected", "OpenAIException - rejected")] + fn the_message_names_the_provider( + #[case] provider: &str, + #[case] body: &str, + #[case] message: &str, + ) { + assert_eq!( + mapped(provider, &http(400, body)), + with_debug(failure( + kind(StatusClass::BadRequest, 400, body), + message, + provider + )) + ); + } +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs new file mode 100644 index 00000000000..d64868c1606 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs @@ -0,0 +1,49 @@ +use super::public::StatusClass; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LocalClass { + ValueError, + FileNotFound, + OsError, +} + +/// A route failure in the shape Python's `exception_type` receives it, before any public +/// class is chosen. +#[derive(Clone, Debug, PartialEq)] +pub enum OriginalException { + Http { + status: u16, + body: String, + headers: Vec<(String, String)>, + }, + Connection { + message: String, + }, + Timeout { + timeout_seconds: Option, + elapsed_seconds: Option, + }, + Response { + message: String, + }, + Local { + class: LocalClass, + message: String, + }, + /// A failure Python raises as a public LiteLLM exception itself, which `exception_type` + /// hands back unchanged. + Public { + class: StatusClass, + message: String, + }, +} + +/// Which of the provider-specific mappers in `exception_type` a route's provider uses. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ExceptionFamily { + OpenAiCompatible, + VertexAi, + Cohere, + #[default] + Other, +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/public.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/public.rs new file mode 100644 index 00000000000..a7319c1287b --- /dev/null +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/public.rs @@ -0,0 +1,262 @@ +use serde::Serialize; + +/// The public LiteLLM classes built from a status code alone: every one takes the same +/// constructor arguments. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::EnumIter, strum::IntoStaticStr)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum StatusClass { + BadRequest, + Authentication, + PermissionDenied, + NotFound, + RateLimit, + ContextWindowExceeded, + ContentPolicyViolation, + InternalServer, + BadGateway, + ServiceUnavailable, + UnsupportedParams, +} + +impl StatusClass { + /// The `status_code` the Python class sets on itself. + pub const fn status_code(self) -> u16 { + match self { + Self::BadRequest + | Self::ContextWindowExceeded + | Self::ContentPolicyViolation + | Self::UnsupportedParams => 400, + Self::Authentication => 401, + Self::PermissionDenied => 403, + Self::NotFound => 404, + Self::RateLimit => 429, + Self::InternalServer => 500, + Self::BadGateway => 502, + Self::ServiceUnavailable => 503, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct UpstreamResponse { + pub status: u16, + pub body: String, + pub headers: Vec<(String, String)>, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct HttpStub { + pub status: u16, + pub method: &'static str, + pub url: &'static str, + pub content: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ResponseArg { + Upstream(UpstreamResponse), + Stub(HttpStub), +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum PublicKind { + Status { + status_class: StatusClass, + response: Option, + }, + Timeout { + status: Option, + }, + ApiConnection, + Api { + status: u16, + request_url: &'static str, + }, +} + +/// Constructor arguments for the public LiteLLM exception, as `exception_type` passes them. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct PublicFailure { + pub kind: PublicKind, + pub message: String, + pub model: String, + pub llm_provider: Option, + pub litellm_debug_info: Option, + pub litellm_response_headers: Option>, + pub print_banner: bool, +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + use std::path::PathBuf; + + use serde_json::Value; + use strum::IntoEnumIterator; + + use super::*; + + const REGENERATE: &str = "LITELLM_REGENERATE_PUBLIC_FAILURE_FIXTURES"; + + fn fixture_directory() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../../tests/test_litellm/rust_bridge/fixtures/public_failures") + } + + fn upstream(status: u16) -> ResponseArg { + ResponseArg::Upstream(UpstreamResponse { + status, + body: r#"{"message": "rejected"}"#.into(), + headers: vec![("retry-after".into(), "7".into())], + }) + } + + fn status_response(class: StatusClass) -> Option { + match class { + StatusClass::Authentication => None, + StatusClass::PermissionDenied => Some(ResponseArg::Stub(HttpStub { + status: 403, + method: "POST", + url: " https://cloud.google.com/vertex-ai/", + content: None, + })), + StatusClass::InternalServer => Some(ResponseArg::Stub(HttpStub { + status: 500, + method: "completion", + url: "https://github.com/BerriAI/litellm", + content: Some("upstream text".into()), + })), + class => Some(upstream(class.status_code())), + } + } + + fn failure(kind: PublicKind, name: &str) -> PublicFailure { + let headers = matches!( + &kind, + PublicKind::Status { + response: Some(ResponseArg::Upstream(_)), + .. + } + ); + PublicFailure { + kind, + message: format!("MistralException - {name}"), + model: "ocr-model".into(), + llm_provider: Some("mistral".into()), + litellm_debug_info: Some("\nModel: ocr-model".into()), + litellm_response_headers: headers.then(|| vec![("retry-after".into(), "7".into())]), + print_banner: false, + } + } + + /// One payload per public class the constructor can build; `test_failures.py` reads the + /// same files, so a shape change on either side fails there or here. + fn fixtures() -> Vec<(String, PublicFailure)> { + let statuses = StatusClass::iter().map(|class| { + let name: &'static str = class.into(); + let name = format!("status_{name}"); + let built = failure( + PublicKind::Status { + status_class: class, + response: status_response(class), + }, + &name, + ); + (name, built) + }); + let others = [ + ( + "timeout_with_status", + PublicFailure { + print_banner: true, + ..failure( + PublicKind::Timeout { status: Some(504) }, + "timeout_with_status", + ) + }, + ), + ( + "timeout_without_status", + PublicFailure { + litellm_debug_info: None, + ..failure( + PublicKind::Timeout { status: None }, + "timeout_without_status", + ) + }, + ), + ( + "api_connection", + PublicFailure { + llm_provider: None, + ..failure(PublicKind::ApiConnection, "api_connection") + }, + ), + ( + "api", + failure( + PublicKind::Api { + status: 409, + request_url: "https://docs.litellm.ai/docs", + }, + "api", + ), + ), + ] + .map(|(name, built)| (name.to_string(), built)); + statuses.chain(others).collect() + } + + #[test] + fn serialized_payloads_match_the_golden_fixtures_python_reads() { + let directory = fixture_directory(); + let regenerate = std::env::var_os(REGENERATE).is_some(); + let expected = fixtures(); + for (name, built) in &expected { + let path = directory.join(format!("{name}.json")); + let serialized = serde_json::to_value(built).unwrap(); + if regenerate { + std::fs::create_dir_all(&directory).unwrap(); + std::fs::write( + &path, + format!("{}\n", serde_json::to_string_pretty(&serialized).unwrap()), + ) + .unwrap(); + } + let golden: Value = + serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); + assert_eq!(serialized, golden, "{name}; set {REGENERATE}=1 to rewrite"); + } + let on_disk: BTreeSet = std::fs::read_dir(&directory) + .unwrap() + .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned()) + .collect(); + let generated: BTreeSet = expected + .iter() + .map(|(name, _)| format!("{name}.json")) + .collect(); + assert_eq!(on_disk, generated); + } + + #[rstest::rstest] + #[case(StatusClass::BadRequest, 400)] + #[case(StatusClass::Authentication, 401)] + #[case(StatusClass::PermissionDenied, 403)] + #[case(StatusClass::NotFound, 404)] + #[case(StatusClass::RateLimit, 429)] + #[case(StatusClass::ContextWindowExceeded, 400)] + #[case(StatusClass::ContentPolicyViolation, 400)] + #[case(StatusClass::InternalServer, 500)] + #[case(StatusClass::BadGateway, 502)] + #[case(StatusClass::ServiceUnavailable, 503)] + #[case(StatusClass::UnsupportedParams, 400)] + fn status_codes_are_the_ones_the_python_classes_set( + #[case] class: StatusClass, + #[case] status: u16, + ) { + assert_eq!(class.status_code(), status); + } +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs new file mode 100644 index 00000000000..5df8ad991b5 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs @@ -0,0 +1,403 @@ +use std::sync::LazyLock; + +use fancy_regex::Regex; +use serde_json::Value; + +use super::Mapping; +use super::public::{HttpStub, PublicFailure, PublicKind, ResponseArg, StatusClass}; + +const GITHUB_URL: &str = "https://github.com/BerriAI/litellm"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ResponseChoice { + Omitted, + Provider, + Stub { status: u16, url: &'static str }, + InternalServerStub, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ApiStatus { + Fixed(u16), + Original, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum Kind { + Status { + class: StatusClass, + response: ResponseChoice, + }, + Timeout(Option), + ApiConnection, + Api { + status: ApiStatus, + request_url: &'static str, + }, +} + +/// One branch of a Python `_map_*_exception` function: when it applies, the class it +/// raises, the message it builds, and whether it passes `litellm_debug_info`. +pub(super) struct Rule { + pub(super) when: fn(&Mapping<'_>) -> bool, + pub(super) kind: Kind, + pub(super) message: fn(&Mapping<'_>) -> String, + pub(super) debug: bool, +} + +/// The first rule that applies decides the failure, as the `if`/`elif` chain does in Python. +pub(super) fn apply(rules: &[Rule], mapping: &Mapping<'_>) -> Option { + rules + .iter() + .find(|rule| (rule.when)(mapping)) + .map(|rule| rule.build(mapping)) +} + +impl Rule { + fn build(&self, mapping: &Mapping<'_>) -> PublicFailure { + let kind = match self.kind { + Kind::Status { class, response } => PublicKind::Status { + status_class: class, + response: response.resolve(mapping), + }, + Kind::Timeout(status) => PublicKind::Timeout { status }, + Kind::ApiConnection => PublicKind::ApiConnection, + Kind::Api { + status, + request_url, + } => PublicKind::Api { + status: match status { + ApiStatus::Fixed(status) => status, + ApiStatus::Original => mapping.original.status.unwrap_or(500), + }, + request_url, + }, + }; + PublicFailure { + kind, + message: (self.message)(mapping), + model: mapping.context.model.clone(), + llm_provider: mapping.context.custom_llm_provider.clone(), + litellm_debug_info: self.debug.then(|| mapping.extra_information.clone()), + litellm_response_headers: None, + print_banner: false, + } + } +} + +impl ResponseChoice { + fn resolve(self, mapping: &Mapping<'_>) -> Option { + match self { + Self::Omitted => None, + Self::Provider => mapping.original.response.clone().map(ResponseArg::Upstream), + Self::Stub { status, url } => Some(ResponseArg::Stub(HttpStub { + status, + method: "POST", + url, + content: None, + })), + Self::InternalServerStub => Some(ResponseArg::Stub(HttpStub { + status: 500, + method: "completion", + url: GITHUB_URL, + content: Some(mapping.original.message.clone()), + })), + } + } +} + +pub(super) fn contains_any(text: &str, markers: &[&str]) -> bool { + markers.iter().any(|marker| text.contains(marker)) +} + +static STANDALONE_429: LazyLock = + LazyLock::new(|| Regex::new(r"\b429\b").expect("valid regex")); +static RATE_LIMIT_PHRASE: LazyLock = + LazyLock::new(|| Regex::new(r"rate[\s_\-]*limit").expect("valid regex")); + +/// `ExceptionCheckers.is_error_str_rate_limit`. +pub(super) fn is_rate_limit(error_str: &str, status: Option) -> bool { + if STANDALONE_429.is_match(error_str).unwrap_or(false) && status == Some(429) { + return true; + } + let lower = error_str.to_lowercase(); + RATE_LIMIT_PHRASE.is_match(&lower).unwrap_or(false) + || lower.contains("service tier capacity exceeded") +} + +/// `ExceptionCheckers.is_error_str_context_window_exceeded`. +pub(super) fn is_context_window_exceeded(error_str: &str) -> bool { + let lower = error_str.to_lowercase(); + if lower.contains("string_above_max_length") { + return false; + } + if lower.contains("invalid 'user'") && lower.contains("string too long") { + return false; + } + contains_any( + &lower, + &[ + "exceed context limit", + "this model's maximum context length is", + "string too long. expected a string with maximum length", + "model's maximum context limit", + "is longer than the model's context length", + "input tokens exceed the configured limit", + "`inputs` tokens + `max_new_tokens` must be", + "exceeds the available context size", + "exceeds the maximum number of tokens allowed", + ], + ) || (lower.contains("current length is") && lower.contains("while limit is")) + || (lower.contains("maximum input length is") && lower.contains("tokens")) +} + +/// The integer `error.code` of a JSON error body, read the way Python's `int()` would. +pub(super) fn body_error_code(error_str: &str) -> Option { + let body: Value = serde_json::from_str(error_str).ok()?; + let Some(Value::Object(error)) = body.as_object()?.get("error") else { + return None; + }; + match error.get("code")? { + Value::Number(number) => number + .as_i64() + .or_else(|| number.as_f64().map(|value| value.trunc() as i64)), + Value::String(code) => code.trim().replace('_', "").parse().ok(), + Value::Bool(flag) => Some(i64::from(*flag)), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::super::testing::{context, failure, http}; + use super::super::{ExceptionFamily, OriginalException, UpstreamResponse}; + use super::*; + + fn first_marker(mapping: &Mapping<'_>) -> bool { + mapping.error_str.contains("first") + } + + fn always(_: &Mapping<'_>) -> bool { + true + } + + fn text(mapping: &Mapping<'_>) -> String { + format!("seen {}", mapping.error_str) + } + + const ORDERED: &[Rule] = &[ + Rule { + when: first_marker, + kind: Kind::Status { + class: StatusClass::NotFound, + response: ResponseChoice::Omitted, + }, + message: text, + debug: false, + }, + Rule { + when: always, + kind: Kind::ApiConnection, + message: text, + debug: true, + }, + ]; + + fn apply_one(kind: Kind, debug: bool, original: &OriginalException) -> Option { + let context = context("mistral", ExceptionFamily::OpenAiCompatible); + let mapping = Mapping::new(&context, original); + apply( + &[Rule { + when: always, + kind, + message: text, + debug, + }], + &mapping, + ) + } + + #[rstest::rstest] + #[case::earlier_rule_wins("first and second", failure( + PublicKind::Status { status_class: StatusClass::NotFound, response: None }, + "seen first and second", + "mistral", + ))] + #[case::later_rule_when_the_earlier_does_not_apply("second", PublicFailure { + litellm_debug_info: Some("\nModel: ocr-model\nMessages: `None`".into()), + ..failure(PublicKind::ApiConnection, "seen second", "mistral") + })] + fn the_first_applicable_rule_decides(#[case] body: &str, #[case] expected: PublicFailure) { + let context = context("mistral", ExceptionFamily::OpenAiCompatible); + let original = http(400, body); + assert_eq!( + apply(ORDERED, &Mapping::new(&context, &original)), + Some(expected) + ); + } + + #[test] + fn no_applicable_rule_leaves_the_failure_to_the_caller() { + let context = context("mistral", ExceptionFamily::OpenAiCompatible); + let original = http(400, "second"); + assert_eq!( + apply(&ORDERED[..1], &Mapping::new(&context, &original)), + None + ); + } + + #[rstest::rstest] + #[case::omitted(ResponseChoice::Omitted, None)] + #[case::provider(ResponseChoice::Provider, Some(ResponseArg::Upstream(UpstreamResponse { + status: 400, + body: "body".into(), + headers: vec![("retry-after".into(), "7".into())], + })))] + #[case::stub( + ResponseChoice::Stub { status: 429, url: "https://stub.test" }, + Some(ResponseArg::Stub(HttpStub { status: 429, method: "POST", url: "https://stub.test", content: None })) + )] + #[case::internal_server_stub( + ResponseChoice::InternalServerStub, + Some(ResponseArg::Stub(HttpStub { + status: 500, + method: "completion", + url: GITHUB_URL, + content: Some("body".into()), + })) + )] + fn response_choices_resolve_against_the_original( + #[case] response: ResponseChoice, + #[case] expected: Option, + ) { + let built = apply_one( + Kind::Status { + class: StatusClass::BadRequest, + response, + }, + false, + &http(400, "body"), + ) + .unwrap(); + assert_eq!( + built.kind, + PublicKind::Status { + status_class: StatusClass::BadRequest, + response: expected, + } + ); + } + + #[rstest::rstest] + #[case::fixed(ApiStatus::Fixed(500), http(409, "body"), 500)] + #[case::original(ApiStatus::Original, http(409, "body"), 409)] + #[case::original_without_a_status( + ApiStatus::Original, + OriginalException::Response { message: "body".into() }, + 500 + )] + fn api_status_is_fixed_or_the_originals( + #[case] status: ApiStatus, + #[case] original: OriginalException, + #[case] expected: u16, + ) { + let built = apply_one( + Kind::Api { + status, + request_url: "https://api.test", + }, + false, + &original, + ) + .unwrap(); + assert_eq!( + built, + failure( + PublicKind::Api { + status: expected, + request_url: "https://api.test" + }, + "seen body", + "mistral" + ) + ); + } + + #[rstest::rstest] + #[case::with_debug(true, Some("\nModel: ocr-model\nMessages: `None`"))] + #[case::without_debug(false, None)] + fn debug_rules_carry_the_extra_information( + #[case] debug: bool, + #[case] expected: Option<&str>, + ) { + let built = apply_one(Kind::Timeout(Some(504)), debug, &http(504, "body")).unwrap(); + assert_eq!( + built, + PublicFailure { + litellm_debug_info: expected.map(str::to_string), + ..failure( + PublicKind::Timeout { status: Some(504) }, + "seen body", + "mistral" + ) + } + ); + } + + #[rstest::rstest] + #[case::standalone_429_with_429_status("got 429 back", Some(429), true)] + #[case::standalone_429_with_other_status("got 429 back", Some(400), false)] + #[case::embedded_429("token4290", Some(429), false)] + #[case::phrase_spaced("Rate Limit reached", None, true)] + #[case::phrase_underscored("rate_limit", None, true)] + #[case::phrase_hyphenated("rate-limit", None, true)] + #[case::service_tier("Service tier capacity exceeded", None, true)] + #[case::unrelated("rejected", Some(429), false)] + fn rate_limit_detection( + #[case] text: &str, + #[case] status: Option, + #[case] expected: bool, + ) { + assert_eq!(is_rate_limit(text, status), expected); + } + + #[rstest::rstest] + #[case::exceed_context_limit("Exceed context limit", true)] + #[case::maximum_context_length("This model's maximum context length is 10", true)] + #[case::string_too_long("string too long. Expected a string with maximum length 5", true)] + #[case::maximum_context_limit("the model's maximum context limit", true)] + #[case::longer_than_context("prompt is longer than the model's context length", true)] + #[case::configured_limit("input tokens exceed the configured limit", true)] + #[case::max_new_tokens("`inputs` tokens + `max_new_tokens` must be <= 10", true)] + #[case::available_context("exceeds the available context size", true)] + #[case::maximum_tokens("exceeds the maximum number of tokens allowed", true)] + #[case::current_and_limit("current length is 9 while limit is 8", true)] + #[case::current_without_limit("current length is 9", false)] + #[case::maximum_input_tokens("maximum input length is 8 tokens", true)] + #[case::maximum_input_without_tokens("maximum input length is 8", false)] + #[case::string_above_max_length_wins("string_above_max_length exceed context limit", false)] + #[case::user_field_is_not_context( + "invalid 'user': string too long. expected a string with maximum length", + false + )] + #[case::unrelated("rejected", false)] + fn context_window_detection(#[case] text: &str, #[case] expected: bool) { + assert_eq!(is_context_window_exceeded(text), expected); + } + + #[rstest::rstest] + #[case::integer(r#"{"error": {"code": 429}}"#, Some(429))] + #[case::float(r#"{"error": {"code": 429.9}}"#, Some(429))] + #[case::string(r#"{"error": {"code": " 4_29 "}}"#, Some(429))] + #[case::boolean(r#"{"error": {"code": true}}"#, Some(1))] + #[case::unparseable_string(r#"{"error": {"code": "slow"}}"#, None)] + #[case::null(r#"{"error": {"code": null}}"#, None)] + #[case::no_code(r#"{"error": {}}"#, None)] + #[case::error_not_an_object(r#"{"error": "429"}"#, None)] + #[case::no_error(r#"{"code": 429}"#, None)] + #[case::not_an_object("[429]", None)] + #[case::not_json("429", None)] + fn body_error_code_reads_the_nested_code(#[case] body: &str, #[case] expected: Option) { + assert_eq!(body_error_code(body), expected); + } +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/status.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/status.rs new file mode 100644 index 00000000000..3d817fe9ae2 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/status.rs @@ -0,0 +1,153 @@ +use super::public::{PublicFailure, StatusClass}; +use super::rules::{ApiStatus, Kind, ResponseChoice, Rule, apply}; +use super::{DOCS_URL, Mapping}; + +const fn with_response(class: StatusClass) -> Kind { + Kind::Status { + class, + response: ResponseChoice::Provider, + } +} + +fn message(mapping: &Mapping<'_>) -> String { + format!("{} - {}", mapping.exception_provider, mapping.error_str) +} + +fn status(mapping: &Mapping<'_>) -> u16 { + mapping.original.status.unwrap_or_default() +} + +/// `_map_exception_by_status`, the fallback for a provider error no provider mapper claimed. +const RULES: &[Rule] = &[ + Rule { + when: |mapping| status(mapping) == 401, + kind: with_response(StatusClass::Authentication), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) == 403, + kind: with_response(StatusClass::PermissionDenied), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) == 404, + kind: with_response(StatusClass::NotFound), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) == 408, + kind: Kind::Timeout(None), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) == 429, + kind: with_response(StatusClass::RateLimit), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) == 500, + kind: with_response(StatusClass::InternalServer), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) == 502, + kind: with_response(StatusClass::BadGateway), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) == 503, + kind: with_response(StatusClass::ServiceUnavailable), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) == 504, + kind: Kind::Timeout(Some(504)), + message, + debug: true, + }, + Rule { + when: |mapping| status(mapping) < 500, + kind: with_response(StatusClass::BadRequest), + message, + debug: true, + }, + Rule { + when: |_| true, + kind: Kind::Api { + status: ApiStatus::Original, + request_url: DOCS_URL, + }, + message, + debug: true, + }, +]; + +/// Only a real provider status of 400 or more reaches the table; a status the HTTP handler +/// synthesized for a failure without a response does not. +pub(super) fn map(mapping: &Mapping<'_>) -> Option { + let status = mapping.original.status?; + if status < 400 || mapping.original.status_is_synthesized { + return None; + } + apply(RULES, mapping) +} + +#[cfg(test)] +mod tests { + use super::super::testing::{context, failure, http, upstream, with_debug}; + use super::super::{ExceptionFamily, OriginalException, PublicKind}; + use super::*; + + fn mapped(original: &OriginalException) -> Option { + let context = context("reducto", ExceptionFamily::Other); + map(&Mapping::new(&context, original)) + } + + fn classified(class: StatusClass, status_code: u16) -> PublicKind { + PublicKind::Status { + status_class: class, + response: upstream(status_code, "rejected"), + } + } + + #[rstest::rstest] + #[case::authentication(401, classified(StatusClass::Authentication, 401))] + #[case::permission_denied(403, classified(StatusClass::PermissionDenied, 403))] + #[case::not_found(404, classified(StatusClass::NotFound, 404))] + #[case::request_timeout(408, PublicKind::Timeout { status: None })] + #[case::rate_limited(429, classified(StatusClass::RateLimit, 429))] + #[case::internal_server(500, classified(StatusClass::InternalServer, 500))] + #[case::bad_gateway(502, classified(StatusClass::BadGateway, 502))] + #[case::service_unavailable(503, classified(StatusClass::ServiceUnavailable, 503))] + #[case::gateway_timeout(504, PublicKind::Timeout { status: Some(504) })] + #[case::lowest_client_error(400, classified(StatusClass::BadRequest, 400))] + #[case::other_client_error(409, classified(StatusClass::BadRequest, 409))] + #[case::highest_client_error(499, classified(StatusClass::BadRequest, 499))] + #[case::other_server_error(501, PublicKind::Api { status: 501, request_url: DOCS_URL })] + fn every_mapped_status_and_the_fallback(#[case] status_code: u16, #[case] kind: PublicKind) { + assert_eq!( + mapped(&http(status_code, "rejected")), + Some(with_debug(failure( + kind, + "ReductoException - rejected", + "reducto" + ))) + ); + } + + #[rstest::rstest] + #[case::below_client_errors(http(399, "rejected"))] + #[case::synthesized(OriginalException::Connection { message: "refused".into() })] + #[case::no_status(OriginalException::Response { message: "bad body".into() })] + fn failures_the_table_does_not_claim(#[case] original: OriginalException) { + assert_eq!(mapped(&original), None); + } +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/vertex_ai.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/vertex_ai.rs new file mode 100644 index 00000000000..0fa8c19c4f6 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/vertex_ai.rs @@ -0,0 +1,574 @@ +use super::public::{PublicFailure, PublicKind, ResponseArg, StatusClass, UpstreamResponse}; +use super::rules::{ + Kind, ResponseChoice, Rule, apply, body_error_code, contains_any, is_context_window_exceeded, +}; +use super::{Mapping, python_capitalize}; + +const VERTEX_URL: &str = "https://cloud.google.com/vertex-ai/"; +const VERTEX_URL_WITH_SPACE: &str = " https://cloud.google.com/vertex-ai/"; + +const QUOTA_MARKERS: &[&str] = &[ + "429 Quota exceeded", + "Quota exceeded for", + "Resource exhausted", + "IndexError: list index out of range", + "429 Unable to submit request because the service is temporarily out of capacity.", +]; + +const fn stubbed(class: StatusClass, status: u16, url: &'static str) -> Kind { + Kind::Status { + class, + response: ResponseChoice::Stub { status, url }, + } +} + +const fn bare(class: StatusClass) -> Kind { + Kind::Status { + class, + response: ResponseChoice::Omitted, + } +} + +/// `{Provider}Exception{label} - {error_str}` with Python's `str.capitalize()`. +fn capitalized(mapping: &Mapping<'_>, label: &str) -> String { + format!( + "{}Exception{label} - {}", + python_capitalize(mapping.provider), + mapping.error_str + ) +} + +/// `litellm.{Class}: {provider}Exception - {error_str}` with the provider as given. +fn litellm_prefixed(mapping: &Mapping<'_>, class: &str) -> String { + format!( + "litellm.{class}: {}Exception - {}", + mapping.provider, mapping.error_str + ) +} + +fn status_is(mapping: &Mapping<'_>, status: u16) -> bool { + mapping.original.status == Some(status) +} + +/// `_map_vertex_exception`, in its branch order. A failure no rule claims falls through +/// to the status table. +const RULES: &[Rule] = &[ + Rule { + when: |mapping| { + contains_any( + &mapping.error_str, + &[ + "Vertex AI API has not been used in project", + "Unable to find your project", + ], + ) + }, + kind: stubbed(StatusClass::BadRequest, 400, VERTEX_URL_WITH_SPACE), + message: |mapping| litellm_prefixed(mapping, "BadRequestError"), + debug: true, + }, + Rule { + when: |mapping| { + mapping + .error_str + .contains("400 Request payload size exceeds") + }, + kind: bare(StatusClass::ContextWindowExceeded), + message: |mapping| capitalized(mapping, ""), + debug: false, + }, + Rule { + when: |mapping| is_context_window_exceeded(&mapping.error_str), + kind: bare(StatusClass::ContextWindowExceeded), + message: |mapping| format!("ContextWindowExceededError: {}", capitalized(mapping, "")), + debug: true, + }, + Rule { + when: |mapping| { + contains_any( + &mapping.error_str, + &["None Unknown Error.", "Content has no parts."], + ) + }, + kind: Kind::Status { + class: StatusClass::InternalServer, + response: ResponseChoice::InternalServerStub, + }, + message: |mapping| litellm_prefixed(mapping, "InternalServerError"), + debug: true, + }, + Rule { + when: |mapping| mapping.error_str.contains("API key not valid."), + kind: bare(StatusClass::Authentication), + message: |mapping| capitalized(mapping, ""), + debug: true, + }, + Rule { + when: |mapping| mapping.error_str.contains("403"), + kind: stubbed(StatusClass::BadRequest, 403, VERTEX_URL_WITH_SPACE), + message: |mapping| capitalized(mapping, " BadRequestError"), + debug: true, + }, + Rule { + when: |mapping| { + contains_any( + &mapping.error_str, + &[ + "The response was blocked.", + "Output blocked by content filtering policy", + ], + ) + }, + kind: stubbed( + StatusClass::ContentPolicyViolation, + 400, + VERTEX_URL_WITH_SPACE, + ), + message: |mapping| capitalized(mapping, " ContentPolicyViolationError"), + debug: true, + }, + Rule { + when: |mapping| { + contains_any(&mapping.error_str, QUOTA_MARKERS) + || (mapping + .original + .status + .is_some_and(|status| (500..600).contains(&status)) + && body_error_code(&mapping.error_str) == Some(429)) + }, + kind: stubbed(StatusClass::RateLimit, 429, VERTEX_URL_WITH_SPACE), + message: |mapping| litellm_prefixed(mapping, "RateLimitError"), + debug: true, + }, + Rule { + when: |mapping| { + contains_any( + &mapping.error_str, + &["500 Internal Server Error", "The model is overloaded."], + ) + }, + kind: bare(StatusClass::InternalServer), + message: |mapping| litellm_prefixed(mapping, "InternalServerError"), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, 400), + kind: stubbed(StatusClass::BadRequest, 400, VERTEX_URL), + message: |mapping| capitalized(mapping, " BadRequestError"), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, 401), + kind: bare(StatusClass::Authentication), + message: |mapping| capitalized(mapping, ""), + debug: false, + }, + Rule { + when: |mapping| status_is(mapping, 403), + kind: stubbed(StatusClass::PermissionDenied, 403, VERTEX_URL), + message: |mapping| capitalized(mapping, ""), + debug: false, + }, + Rule { + when: |mapping| status_is(mapping, 404), + kind: bare(StatusClass::NotFound), + message: |mapping| capitalized(mapping, ""), + debug: false, + }, + Rule { + when: |mapping| status_is(mapping, 408), + kind: Kind::Timeout(None), + message: |mapping| capitalized(mapping, ""), + debug: false, + }, + Rule { + when: |mapping| status_is(mapping, 429), + kind: stubbed(StatusClass::RateLimit, 429, VERTEX_URL_WITH_SPACE), + message: |mapping| format!("litellm.RateLimitError: {}", capitalized(mapping, "")), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, 500), + kind: Kind::Status { + class: StatusClass::InternalServer, + response: ResponseChoice::InternalServerStub, + }, + message: |mapping| capitalized(mapping, " InternalServerError"), + debug: true, + }, + Rule { + when: |mapping| status_is(mapping, 502), + kind: Kind::ApiConnection, + message: |mapping| capitalized(mapping, ""), + debug: false, + }, + Rule { + when: |mapping| status_is(mapping, 503), + kind: bare(StatusClass::ServiceUnavailable), + message: |mapping| capitalized(mapping, ""), + debug: false, + }, +]; + +pub(super) fn map(mapping: &Mapping<'_>) -> Option { + apply(RULES, mapping).map(|failure| keep_upstream_response(mapping, failure)) +} + +/// Deliberate divergence from `_map_vertex_exception`, which replaces the provider response +/// with a stub and so drops the upstream body and `retry-after`. The response keeps the +/// status the public class carries. +fn keep_upstream_response(mapping: &Mapping<'_>, failure: PublicFailure) -> PublicFailure { + let (PublicKind::Status { status_class, .. }, Some(upstream), false) = ( + &failure.kind, + &mapping.original.response, + mapping.original.status_is_synthesized, + ) else { + return failure; + }; + PublicFailure { + kind: PublicKind::Status { + status_class: *status_class, + response: Some(ResponseArg::Upstream(UpstreamResponse { + status: status_class.status_code(), + ..upstream.clone() + })), + }, + ..failure + } +} + +#[cfg(test)] +mod tests { + use super::super::testing::{context, failure, http, status, upstream, with_debug}; + use super::super::{ExceptionFamily, HttpStub, OriginalException}; + use super::*; + + fn mapped(original: &OriginalException) -> Option { + let context = context("vertex_ai", ExceptionFamily::VertexAi); + map(&Mapping::new(&context, original)) + } + + fn kept(class: StatusClass, body: &str) -> PublicKind { + status(class, upstream(class.status_code(), body)) + } + + #[rstest::rstest] + #[case::api_not_enabled( + 400, + "Vertex AI API has not been used in project x", + with_debug(failure( + kept( + StatusClass::BadRequest, + "Vertex AI API has not been used in project x" + ), + "litellm.BadRequestError: vertex_aiException - Vertex AI API has not been used in project x", + "vertex_ai", + )) + )] + #[case::project_not_found( + 400, + "Unable to find your project", + with_debug(failure( + kept(StatusClass::BadRequest, "Unable to find your project"), + "litellm.BadRequestError: vertex_aiException - Unable to find your project", + "vertex_ai", + )) + )] + #[case::payload_too_large( + 400, + "400 Request payload size exceeds the limit", + failure( + kept( + StatusClass::ContextWindowExceeded, + "400 Request payload size exceeds the limit" + ), + "Vertex_aiException - 400 Request payload size exceeds the limit", + "vertex_ai", + ) + )] + #[case::context_window( + 500, + "This model's maximum context length is 10", + with_debug(failure( + kept( + StatusClass::ContextWindowExceeded, + "This model's maximum context length is 10" + ), + "ContextWindowExceededError: Vertex_aiException - This model's maximum context length is 10", + "vertex_ai", + )) + )] + #[case::unknown_error( + 400, + "None Unknown Error.", + with_debug(failure( + kept(StatusClass::InternalServer, "None Unknown Error."), + "litellm.InternalServerError: vertex_aiException - None Unknown Error.", + "vertex_ai", + )) + )] + #[case::no_parts( + 400, + "Content has no parts.", + with_debug(failure( + kept(StatusClass::InternalServer, "Content has no parts."), + "litellm.InternalServerError: vertex_aiException - Content has no parts.", + "vertex_ai", + )) + )] + #[case::api_key_not_valid( + 400, + "API key not valid.", + with_debug(failure( + kept(StatusClass::Authentication, "API key not valid."), + "Vertex_aiException - API key not valid.", + "vertex_ai", + )) + )] + #[case::forbidden_text( + 400, + "got a 403", + with_debug(failure( + kept(StatusClass::BadRequest, "got a 403"), + "Vertex_aiException BadRequestError - got a 403", + "vertex_ai", + )) + )] + #[case::response_blocked( + 400, + "The response was blocked.", + with_debug(failure( + kept(StatusClass::ContentPolicyViolation, "The response was blocked."), + "Vertex_aiException ContentPolicyViolationError - The response was blocked.", + "vertex_ai", + )) + )] + #[case::output_blocked( + 400, + "Output blocked by content filtering policy", + with_debug(failure( + kept( + StatusClass::ContentPolicyViolation, + "Output blocked by content filtering policy" + ), + "Vertex_aiException ContentPolicyViolationError - Output blocked by content filtering policy", + "vertex_ai", + )) + )] + #[case::quota_marker( + 400, + "Quota exceeded for aiplatform", + with_debug(failure( + kept(StatusClass::RateLimit, "Quota exceeded for aiplatform"), + "litellm.RateLimitError: vertex_aiException - Quota exceeded for aiplatform", + "vertex_ai", + )) + )] + #[case::wrapped_429( + 503, + r#"{"error": {"code": "429"}}"#, + with_debug(failure( + kept(StatusClass::RateLimit, r#"{"error": {"code": "429"}}"#), + r#"litellm.RateLimitError: vertex_aiException - {"error": {"code": "429"}}"#, + "vertex_ai", + )) + )] + #[case::overloaded( + 400, + "The model is overloaded.", + with_debug(failure( + kept(StatusClass::InternalServer, "The model is overloaded."), + "litellm.InternalServerError: vertex_aiException - The model is overloaded.", + "vertex_ai", + )) + )] + #[case::internal_server_text( + 400, + "500 Internal Server Error", + with_debug(failure( + kept(StatusClass::InternalServer, "500 Internal Server Error"), + "litellm.InternalServerError: vertex_aiException - 500 Internal Server Error", + "vertex_ai", + )) + )] + fn each_text_rule_maps_by_the_body( + #[case] status_code: u16, + #[case] body: &str, + #[case] expected: PublicFailure, + ) { + assert_eq!(mapped(&http(status_code, body)), Some(expected)); + } + + #[rstest::rstest] + #[case::bad_request( + 400, + with_debug(failure( + kept(StatusClass::BadRequest, "rejected"), + "Vertex_aiException BadRequestError - rejected", + "vertex_ai" + )) + )] + #[case::authentication( + 401, + failure( + kept(StatusClass::Authentication, "rejected"), + "Vertex_aiException - rejected", + "vertex_ai" + ) + )] + #[case::permission_denied( + 403, + failure( + kept(StatusClass::PermissionDenied, "rejected"), + "Vertex_aiException - rejected", + "vertex_ai" + ) + )] + #[case::not_found( + 404, + failure( + kept(StatusClass::NotFound, "rejected"), + "Vertex_aiException - rejected", + "vertex_ai" + ) + )] + #[case::request_timeout(408, failure(PublicKind::Timeout { status: None }, "Vertex_aiException - rejected", "vertex_ai"))] + #[case::rate_limited( + 429, + with_debug(failure( + kept(StatusClass::RateLimit, "rejected"), + "litellm.RateLimitError: Vertex_aiException - rejected", + "vertex_ai" + )) + )] + #[case::internal_server( + 500, + with_debug(failure( + kept(StatusClass::InternalServer, "rejected"), + "Vertex_aiException InternalServerError - rejected", + "vertex_ai" + )) + )] + #[case::bad_gateway( + 502, + failure( + PublicKind::ApiConnection, + "Vertex_aiException - rejected", + "vertex_ai" + ) + )] + #[case::service_unavailable( + 503, + failure( + kept(StatusClass::ServiceUnavailable, "rejected"), + "Vertex_aiException - rejected", + "vertex_ai" + ) + )] + fn each_status_rule_maps_by_the_status( + #[case] status_code: u16, + #[case] expected: PublicFailure, + ) { + assert_eq!(mapped(&http(status_code, "rejected")), Some(expected)); + } + + #[rstest::rstest] + #[case::unmapped_status(409)] + #[case::gateway_timeout(504)] + fn statuses_without_a_rule_fall_through(#[case] status_code: u16) { + assert_eq!(mapped(&http(status_code, "rejected")), None); + } + + #[rstest::rstest] + #[case::stub_without_an_upstream_response( + OriginalException::Response { message: "got a 403".into() }, + status(StatusClass::BadRequest, Some(ResponseArg::Stub(HttpStub { status: 403, method: "POST", url: VERTEX_URL_WITH_SPACE, content: None }))) + )] + #[case::stub_for_a_synthesized_status( + OriginalException::Connection { message: "got a 403".into() }, + status(StatusClass::BadRequest, Some(ResponseArg::Stub(HttpStub { status: 403, method: "POST", url: VERTEX_URL_WITH_SPACE, content: None }))) + )] + fn the_rule_response_stays_when_there_is_no_real_upstream_response( + #[case] original: OriginalException, + #[case] kind: PublicKind, + ) { + assert_eq!(mapped(&original).map(|failure| failure.kind), Some(kind)); + } + + #[test] + fn a_synthesized_500_keeps_the_internal_server_stub() { + let original = OriginalException::Connection { + message: "refused".into(), + }; + assert_eq!( + mapped(&original), + Some(with_debug(failure( + status( + StatusClass::InternalServer, + Some(ResponseArg::Stub(HttpStub { + status: 500, + method: "completion", + url: "https://github.com/BerriAI/litellm", + content: Some("refused".into()), + })) + ), + "Vertex_aiException InternalServerError - refused", + "vertex_ai" + ))) + ); + } + + #[rstest::rstest] + #[case::project_before_payload_size( + "Unable to find your project 400 Request payload size exceeds", + StatusClass::BadRequest, + "litellm.BadRequestError: vertex_aiException - Unable to find your project 400 Request payload size exceeds", + true + )] + #[case::payload_size_before_context_window( + "400 Request payload size exceeds; This model's maximum context length is 10", + StatusClass::ContextWindowExceeded, + "Vertex_aiException - 400 Request payload size exceeds; This model's maximum context length is 10", + false + )] + #[case::api_key_before_forbidden( + "API key not valid. 403", + StatusClass::Authentication, + "Vertex_aiException - API key not valid. 403", + true + )] + #[case::forbidden_before_blocked( + "403 The response was blocked.", + StatusClass::BadRequest, + "Vertex_aiException BadRequestError - 403 The response was blocked.", + true + )] + #[case::blocked_before_quota( + "The response was blocked. Resource exhausted", + StatusClass::ContentPolicyViolation, + "Vertex_aiException ContentPolicyViolationError - The response was blocked. Resource exhausted", + true + )] + #[case::quota_before_overloaded( + "Resource exhausted The model is overloaded.", + StatusClass::RateLimit, + "litellm.RateLimitError: vertex_aiException - Resource exhausted The model is overloaded.", + true + )] + fn the_earlier_rule_wins_when_two_apply( + #[case] body: &str, + #[case] class: StatusClass, + #[case] message: &str, + #[case] debug: bool, + ) { + let expected = failure(kept(class, body), message, "vertex_ai"); + assert_eq!( + mapped(&http(401, body)), + Some(if debug { + with_debug(expected) + } else { + expected + }) + ); + } +} diff --git a/litellm-rust/crates/core-utils/src/lib.rs b/litellm-rust/crates/core-utils/src/lib.rs index a8895cccf2c..0c26aa50cc3 100644 --- a/litellm-rust/crates/core-utils/src/lib.rs +++ b/litellm-rust/crates/core-utils/src/lib.rs @@ -1,7 +1,10 @@ pub mod call_arguments; pub mod core_helpers; +pub mod exception_mapping_utils; pub mod get_llm_provider_logic; pub mod params; pub mod prompt_templates; +pub mod python_repr; +pub mod secret_redaction; pub mod serde_compat; pub mod url_utils; diff --git a/litellm-rust/crates/core-utils/src/python_repr.rs b/litellm-rust/crates/core-utils/src/python_repr.rs new file mode 100644 index 00000000000..7ccfb826377 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/python_repr.rs @@ -0,0 +1,93 @@ +/// `repr()` of a Python `str`: single quotes unless the text holds a single quote and no +/// double quote, with backslashes, the chosen quote and control characters escaped. +pub fn python_str_repr(value: &str) -> String { + let quote = if value.contains('\'') && !value.contains('"') { + '"' + } else { + '\'' + }; + let escaped: String = value + .chars() + .map(|character| match character { + '\\' => "\\\\".to_string(), + '\t' => "\\t".to_string(), + '\n' => "\\n".to_string(), + '\r' => "\\r".to_string(), + character if character == quote => format!("\\{character}"), + character + if (character as u32) < 0x20 || (0x7f..0xa0).contains(&(character as u32)) => + { + format!("\\x{:02x}", character as u32) + } + character => character.to_string(), + }) + .collect(); + format!("{quote}{escaped}{quote}") +} + +/// `repr()` of the Python value a JSON value decodes to. +pub fn python_value_repr(value: &serde_json::Value) -> String { + use serde_json::Value; + match value { + Value::Null => "None".to_string(), + Value::Bool(true) => "True".to_string(), + Value::Bool(false) => "False".to_string(), + Value::Number(number) => number.to_string(), + Value::String(text) => python_str_repr(text), + Value::Array(items) => format!( + "[{}]", + items + .iter() + .map(python_value_repr) + .collect::>() + .join(", ") + ), + Value::Object(fields) => format!( + "{{{}}}", + fields + .iter() + .map(|(key, value)| format!( + "{}: {}", + python_str_repr(key), + python_value_repr(value) + )) + .collect::>() + .join(", ") + ), + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{python_str_repr, python_value_repr}; + + #[rstest::rstest] + #[case::null(json!(null), "None")] + #[case::true_(json!(true), "True")] + #[case::false_(json!(false), "False")] + #[case::integer(json!(5), "5")] + #[case::float(json!(1.5), "1.5")] + #[case::string(json!("it's"), "\"it's\"")] + #[case::list(json!(["a", 1]), "['a', 1]")] + #[case::dict(json!({"format": "native"}), "{'format': 'native'}")] + #[case::empty_list(json!([]), "[]")] + fn value_repr_matches_python(#[case] value: serde_json::Value, #[case] expected: &str) { + assert_eq!(python_value_repr(&value), expected); + } + + #[rstest::rstest] + #[case::plain("native", "'native'")] + #[case::single_quote("it's", "\"it's\"")] + #[case::both_quotes("it's \"x\"", "'it\\'s \"x\"'")] + #[case::double_quote("say \"x\"", "'say \"x\"'")] + #[case::backslash("a\\b", "'a\\\\b'")] + #[case::whitespace("a\tb\nc\rd", "'a\\tb\\nc\\rd'")] + #[case::control("a\u{1}b\u{7f}c\u{85}", "'a\\x01b\\x7fc\\x85'")] + #[case::unicode("café", "'café'")] + #[case::empty("", "''")] + fn matches_python_repr(#[case] value: &str, #[case] expected: &str) { + assert_eq!(python_str_repr(value), expected); + } +} diff --git a/litellm-rust/crates/core-utils/src/secret_redaction.rs b/litellm-rust/crates/core-utils/src/secret_redaction.rs new file mode 100644 index 00000000000..32922d7430c --- /dev/null +++ b/litellm-rust/crates/core-utils/src/secret_redaction.rs @@ -0,0 +1,97 @@ +use std::sync::LazyLock; + +use fancy_regex::Regex; + +pub const REDACTED: &str = "REDACTED"; + +const DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH: usize = 16; + +fn minimum_custom_key_length() -> usize { + std::env::var("MINIMUM_CUSTOM_KEY_LENGTH") + .ok() + .and_then(|value| value.trim().parse().ok()) + .unwrap_or(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH) +} + +fn secret_patterns(minimum_custom_key_length: usize) -> String { + let sk_suffix_length = minimum_custom_key_length.saturating_sub("sk-".len()); + [ + r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----", + r"\bya29\.[A-Za-z0-9_.~+/-]+", + r#"(?:client_secret|azure_password|azure_username)\s+[^\s,'"})\]{}>]+"#, + r"(?:AKIA|ASIA)[0-9A-Z]{16}", + r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*", + r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}", + &format!(r"sk-[A-Za-z0-9\-_]{{{sk_suffix_length},}}"), + r#"(?<=[?&])(?:api[_-]?key|\w*(?:token|password|passwd|client_secret|secret_key|_secret))=[^\s&'"]+"#, + r#"(?:api[_-]?key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]{8,}"#, + r#"(?:x-api-key|api-key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, + r"x-ak-[A-Za-z0-9\-_]{20,}", + r"AIza[0-9A-Za-z\-_]{35}", + r#"(?<=[?&])key=[^\s&'"]{8,}"#, + r#"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, + r#"(?<=://)[^\s'":]{0,4096}:[^\s'"]{1,4096}(?=@)"#, + r"dapi[0-9a-f]{32}", + r#"litellm\.[A-Za-z0-9_]*_key['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, + r#"private_key['"]?\s*[:=]\s*['"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'"})\]{}>]+)"#, + concat!( + r"(?:master_key|xai_key|database_url|db_url|connection_string|", + r"aws_secret_access_key|aws_session_token|aws_access_key_id|", + r"signing_key|encryption_key|", + r"auth_token|access_token|refresh_token|", + r"slack_webhook_url|webhook_url|", + r"database_connection_string|", + r"huggingface_token|jwt_secret)", + r#"['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, + ), + r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*", + r"(?<=[?&])sig=[A-Za-z0-9%+/=]+", + r#"\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"#, + ] + .join("|") +} + +static SECRET_RE: LazyLock = LazyLock::new(|| { + Regex::new(&format!( + "(?i){}", + secret_patterns(minimum_custom_key_length()) + )) + .expect("secret redaction patterns compile") +}); + +pub fn redact_string(value: &str) -> String { + SECRET_RE.replace_all(value, REDACTED).into_owned() +} + +pub fn secret_redaction_enabled() -> bool { + !std::env::var("LITELLM_DISABLE_REDACT_SECRETS") + .is_ok_and(|value| value.eq_ignore_ascii_case("true")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[rstest::rstest] + #[case::bearer("auth failed: Bearer abcdefghijklmnop", "auth failed: REDACTED")] + #[case::sk_key("key sk-abcdefghijklmnopqrstuvwxyz rejected", "key REDACTED rejected")] + #[case::short_sk_key_is_kept("sk-abc", "sk-abc")] + #[case::query_param("GET /v1?api_key=secret123&x=1", "GET /v1?REDACTED&x=1")] + #[case::dict_repr("{'api_key': 'abcdefghij'}", "{'REDACTED'}")] + #[case::url_credentials("postgres://user:pass@host/db", "postgres://REDACTED@host/db")] + #[case::case_insensitive("BEARER ABCDEFGHIJKLMNOP", "REDACTED")] + #[case::aws_key("AKIAABCDEFGHIJKLMNOP", "REDACTED")] + #[case::sas_signature("https://x.blob/a?sv=1&sig=abc%2B=", "https://x.blob/a?sv=1&REDACTED")] + #[case::password_needs_word_boundary("db_password=hunter2", "REDACTED")] + #[case::plain_text_is_kept(r#"{"message": "rejected"}"#, r#"{"message": "rejected"}"#)] + fn redacts_the_same_spans_as_the_python_patterns(#[case] input: &str, #[case] expected: &str) { + assert_eq!(redact_string(input), expected); + } + + #[test] + fn sk_threshold_follows_the_minimum_custom_key_length() { + let patterns = Regex::new(&format!("(?i){}", secret_patterns(8))).unwrap(); + assert_eq!(patterns.replace_all("sk-abcde", REDACTED), REDACTED); + assert_eq!(patterns.replace_all("sk-abcd", REDACTED), "sk-abcd"); + } +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/api.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/api.json new file mode 100644 index 00000000000..ae629114589 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/api.json @@ -0,0 +1,13 @@ +{ + "kind": { + "type": "api", + "status": 409, + "request_url": "https://docs.litellm.ai/docs" + }, + "message": "MistralException - api", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": null, + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/api_connection.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/api_connection.json new file mode 100644 index 00000000000..ab400ed9e01 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/api_connection.json @@ -0,0 +1,11 @@ +{ + "kind": { + "type": "api_connection" + }, + "message": "MistralException - api_connection", + "model": "ocr-model", + "llm_provider": null, + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": null, + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_authentication.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_authentication.json new file mode 100644 index 00000000000..057392575a1 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_authentication.json @@ -0,0 +1,13 @@ +{ + "kind": { + "type": "status", + "status_class": "authentication", + "response": null + }, + "message": "MistralException - status_authentication", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": null, + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_gateway.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_gateway.json new file mode 100644 index 00000000000..abb1425f686 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_gateway.json @@ -0,0 +1,28 @@ +{ + "kind": { + "type": "status", + "status_class": "bad_gateway", + "response": { + "type": "upstream", + "status": 502, + "body": "{\"message\": \"rejected\"}", + "headers": [ + [ + "retry-after", + "7" + ] + ] + } + }, + "message": "MistralException - status_bad_gateway", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": [ + [ + "retry-after", + "7" + ] + ], + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_request.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_request.json new file mode 100644 index 00000000000..171a994cd35 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_request.json @@ -0,0 +1,28 @@ +{ + "kind": { + "type": "status", + "status_class": "bad_request", + "response": { + "type": "upstream", + "status": 400, + "body": "{\"message\": \"rejected\"}", + "headers": [ + [ + "retry-after", + "7" + ] + ] + } + }, + "message": "MistralException - status_bad_request", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": [ + [ + "retry-after", + "7" + ] + ], + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_content_policy_violation.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_content_policy_violation.json new file mode 100644 index 00000000000..ae9c1e145de --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_content_policy_violation.json @@ -0,0 +1,28 @@ +{ + "kind": { + "type": "status", + "status_class": "content_policy_violation", + "response": { + "type": "upstream", + "status": 400, + "body": "{\"message\": \"rejected\"}", + "headers": [ + [ + "retry-after", + "7" + ] + ] + } + }, + "message": "MistralException - status_content_policy_violation", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": [ + [ + "retry-after", + "7" + ] + ], + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_context_window_exceeded.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_context_window_exceeded.json new file mode 100644 index 00000000000..61e1a56a622 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_context_window_exceeded.json @@ -0,0 +1,28 @@ +{ + "kind": { + "type": "status", + "status_class": "context_window_exceeded", + "response": { + "type": "upstream", + "status": 400, + "body": "{\"message\": \"rejected\"}", + "headers": [ + [ + "retry-after", + "7" + ] + ] + } + }, + "message": "MistralException - status_context_window_exceeded", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": [ + [ + "retry-after", + "7" + ] + ], + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_internal_server.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_internal_server.json new file mode 100644 index 00000000000..b3c5c51a785 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_internal_server.json @@ -0,0 +1,19 @@ +{ + "kind": { + "type": "status", + "status_class": "internal_server", + "response": { + "type": "stub", + "status": 500, + "method": "completion", + "url": "https://github.com/BerriAI/litellm", + "content": "upstream text" + } + }, + "message": "MistralException - status_internal_server", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": null, + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_not_found.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_not_found.json new file mode 100644 index 00000000000..26ffd872961 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_not_found.json @@ -0,0 +1,28 @@ +{ + "kind": { + "type": "status", + "status_class": "not_found", + "response": { + "type": "upstream", + "status": 404, + "body": "{\"message\": \"rejected\"}", + "headers": [ + [ + "retry-after", + "7" + ] + ] + } + }, + "message": "MistralException - status_not_found", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": [ + [ + "retry-after", + "7" + ] + ], + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_permission_denied.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_permission_denied.json new file mode 100644 index 00000000000..42772f98830 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_permission_denied.json @@ -0,0 +1,19 @@ +{ + "kind": { + "type": "status", + "status_class": "permission_denied", + "response": { + "type": "stub", + "status": 403, + "method": "POST", + "url": " https://cloud.google.com/vertex-ai/", + "content": null + } + }, + "message": "MistralException - status_permission_denied", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": null, + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_rate_limit.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_rate_limit.json new file mode 100644 index 00000000000..c9b88822ebd --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_rate_limit.json @@ -0,0 +1,28 @@ +{ + "kind": { + "type": "status", + "status_class": "rate_limit", + "response": { + "type": "upstream", + "status": 429, + "body": "{\"message\": \"rejected\"}", + "headers": [ + [ + "retry-after", + "7" + ] + ] + } + }, + "message": "MistralException - status_rate_limit", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": [ + [ + "retry-after", + "7" + ] + ], + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_service_unavailable.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_service_unavailable.json new file mode 100644 index 00000000000..5af65202ef1 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_service_unavailable.json @@ -0,0 +1,28 @@ +{ + "kind": { + "type": "status", + "status_class": "service_unavailable", + "response": { + "type": "upstream", + "status": 503, + "body": "{\"message\": \"rejected\"}", + "headers": [ + [ + "retry-after", + "7" + ] + ] + } + }, + "message": "MistralException - status_service_unavailable", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": [ + [ + "retry-after", + "7" + ] + ], + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_unsupported_params.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_unsupported_params.json new file mode 100644 index 00000000000..a1b318ce03c --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_unsupported_params.json @@ -0,0 +1,28 @@ +{ + "kind": { + "type": "status", + "status_class": "unsupported_params", + "response": { + "type": "upstream", + "status": 400, + "body": "{\"message\": \"rejected\"}", + "headers": [ + [ + "retry-after", + "7" + ] + ] + } + }, + "message": "MistralException - status_unsupported_params", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": [ + [ + "retry-after", + "7" + ] + ], + "print_banner": false +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_with_status.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_with_status.json new file mode 100644 index 00000000000..8b215bdb7e6 --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_with_status.json @@ -0,0 +1,12 @@ +{ + "kind": { + "type": "timeout", + "status": 504 + }, + "message": "MistralException - timeout_with_status", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": "\nModel: ocr-model", + "litellm_response_headers": null, + "print_banner": true +} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_without_status.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_without_status.json new file mode 100644 index 00000000000..4a79169776e --- /dev/null +++ b/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_without_status.json @@ -0,0 +1,12 @@ +{ + "kind": { + "type": "timeout", + "status": null + }, + "message": "MistralException - timeout_without_status", + "model": "ocr-model", + "llm_provider": "mistral", + "litellm_debug_info": null, + "litellm_response_headers": null, + "print_banner": false +} From 5bbdc72879cbbfd2c3970b09fb189872964484c3 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 20:02:48 +0000 Subject: [PATCH 056/102] fix(gateway): expose /deepgram/ passthrough routes on the gateway allowlist Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- gateway/routes/allowlist.py | 1 + 1 file changed, 1 insertion(+) diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 1c503a083f1..30a5374e91d 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -94,6 +94,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( "/vertex-ai/", "/assemblyai/", "/eu.assemblyai/", + "/deepgram/", "/langfuse/", "/vllm/", "/mistral/", From d37320f8cf242f310bad3e759c9a7dbbdc27b699 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 18 Sep 2026 13:08:55 -0700 Subject: [PATCH 057/102] fix(proxy): report null cost for unpriced deployments on the /model/info id lookup GET /model/info?litellm_model_id= went through _get_proxy_model_info, a copy of _enrich_model_info_with_litellm_data that missed the unpriced check, so the same deployment read as null in the list and as 0 on the id lookup. The id lookup now delegates to the shared helper, so the two paths cannot drift again --- litellm/proxy/proxy_server.py | 40 +-------------- .../proxy_server/test_routes_model_info.py | 50 +++++++++++++++++++ 2 files changed, 51 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f21c907c292..f3b2dbd0b03 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15057,45 +15057,7 @@ def _translate_model_name_for_response(model: dict) -> dict: def _get_proxy_model_info(model: dict) -> dict: - # provided model_info in config.yaml - model_info: Final = model.get("model_info", {}) - - # read litellm model_prices_and_context_window.json to get the following: - # input_cost_per_token, output_cost_per_token, max_tokens - litellm_model_info = get_litellm_model_info(model=model) - - # 2nd pass on the model, try seeing if we can find model in litellm model_cost map - if litellm_model_info == {}: - # use litellm_param model_name to get model_info - litellm_params = model.get("litellm_params", {}) - litellm_model = litellm_params.get("model", None) - try: - litellm_model_info = litellm.get_model_info(model=litellm_model) - except Exception: - litellm_model_info = {} - # 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map - if litellm_model_info == {}: - # use litellm_param model_name to get model_info - litellm_params = model.get("litellm_params", {}) - litellm_model = litellm_params.get("model", None) - split_model: Final = litellm_model.split("/") - if len(split_model) > 0: - litellm_model = split_model[-1] - try: - litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0]) - except Exception: - litellm_model_info = {} - discovered_model_info: Final = ( - llm_router.get_discovered_model_info(model_info.get("id")) if llm_router is not None else MappingProxyType({}) - ) - for k, v in MappingProxyType({**litellm_model_info, **discovered_model_info}).items(): - if k not in model_info or (model_info[k] is None and k in discovered_model_info): - model_info[k] = v - model["model_info"] = model_info - # don't return the llm credentials - model = remove_sensitive_info_from_deployment(deployment_dict=model, excluded_keys={"litellm_credential_name"}) - - return _translate_model_name_for_response(model) + return _translate_model_name_for_response(_enrich_model_info_with_litellm_data(model=model, llm_router=llm_router)) def _model_info_json_response(data: Sequence[Mapping[str, object]] | Mapping[str, object]) -> Response: diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 203a237113f..75a8657356a 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -323,6 +323,56 @@ def test_model_info_reports_null_cost_for_unpriced_deployment_and_zero_for_decla assert input_cost > 0 and output_cost > 0 +def test_model_info_id_lookup_reports_the_same_cost_as_the_list( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``GET /model/info?litellm_model_id=`` must agree with the ``GET /model/info`` list, so an + unpriced deployment cannot read as null in the list and as free on the id lookup.""" + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + declared: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} + free: Final = {"input_cost_per_token": 0, "output_cost_per_token": 0} + router: Final = litellm.Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": f"openai/{name}", "api_key": "x", "api_base": "http://vllm", **costs}, + "model_info": {"id": f"{name}-id"}, + } + for name, costs in (("vllm-unpriced", {}), ("vllm-free", free), ("vllm-priced", declared)) + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + def costs_of(response: httpx.Response) -> dict[str, tuple[object, object]]: + assert response.status_code == 200, response.text + return { + row["model_info"]["id"]: ( + row["model_info"].get("input_cost_per_token"), + row["model_info"].get("output_cost_per_token"), + ) + for row in response.json()["data"] + } + + with auth_as(): + listed: Final = costs_of(client.get("/model/info")) + by_id: Final = { + model_id: costs_of(client.get("/model/info", params={"litellm_model_id": model_id}))[model_id] + for model_id in listed + } + + assert listed == { + "vllm-unpriced-id": (None, None), + "vllm-free-id": (0, 0), + "vllm-priced-id": (declared["input_cost_per_token"], declared["output_cost_per_token"]), + } + assert by_id == listed + _invalidate_model_cost_lowercase_map() + + def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch): from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth from litellm.proxy.auth import model_checks From d1563e0b552e79eea3b8c1a84bbd7808d467cae0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:09:13 -0700 Subject: [PATCH 058/102] feat: honor eager_input_streaming on Bedrock and Anthropic Claude tools --- litellm/llms/anthropic/chat/transformation.py | 19 +++- litellm/llms/anthropic/common_utils.py | 27 ++++- .../adapters/transformation.py | 11 +- .../bedrock/chat/converse_transformation.py | 12 +- .../anthropic_claude3_transformation.py | 12 +- litellm/llms/bedrock/common_utils.py | 14 ++- .../anthropic_claude3_transformation.py | 9 ++ litellm/proxy/_lazy_openapi_snapshot.json | 10 +- litellm/types/llms/anthropic.py | 3 + litellm/types/llms/openai.py | 2 + .../test_anthropic_chat_transformation.py | 71 ++++++++++++ ...al_pass_through_adapters_transformation.py | 42 +++++++ ...ations_anthropic_claude3_transformation.py | 53 +++++++++ .../chat/test_converse_transformation.py | 106 ++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 62 ++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 16 files changed, 444 insertions(+), 13 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 0f99441a115..1f90d375bc2 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -94,6 +94,7 @@ from litellm.utils import ( from ..common_utils import ( AnthropicError, AnthropicModelInfo, + eager_input_streaming_flag, process_anthropic_headers, strip_advisor_blocks_from_messages, ) @@ -732,10 +733,20 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): input_anthropic_schema: Final = sanitize_input_schema_for_anthropic(_input_schema) - _tool: Final = AnthropicMessagesTool( - name=tool["function"]["name"], - input_schema=input_anthropic_schema, - type="custom", + _eager_input_streaming: Final = eager_input_streaming_flag(tool) + _tool: Final = ( + AnthropicMessagesTool( + name=tool["function"]["name"], + input_schema=input_anthropic_schema, + type="custom", + ) + if _eager_input_streaming is None + else AnthropicMessagesTool( + name=tool["function"]["name"], + input_schema=input_anthropic_schema, + type="custom", + eager_input_streaming=_eager_input_streaming, + ) ) _description: Final = tool["function"].get("description") diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d35a9372058..2de9ab41d95 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -10,7 +10,7 @@ from types import MappingProxyType from typing import Any, Final, Literal import httpx -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError import litellm from litellm.constants import ( @@ -19,6 +19,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, ) +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, is_encrypted_reasoning_block, @@ -231,6 +232,27 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup return headers, api_key +class _EagerInputStreamingFunction(BaseModel): + eager_input_streaming: StrictBool | None = None + + +class _EagerInputStreamingTool(BaseModel): + eager_input_streaming: StrictBool | None = None + function: _EagerInputStreamingFunction | None = None + + +def eager_input_streaming_flag(tool: object) -> bool | None: + try: + parsed: Final = _EagerInputStreamingTool.model_validate(tool) + except ValidationError as error: + if isinstance(tool, Mapping): + raise UnsupportedParamsError(message="eager_input_streaming must be a boolean") from error + return None + if parsed.eager_input_streaming is not None: + return parsed.eager_input_streaming + return parsed.function.eager_input_streaming if parsed.function is not None else None + + class AnthropicError(BaseLLMException): def __init__( self, @@ -373,6 +395,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): return False + def is_eager_input_streaming_used(self, tools: Sequence[object] | None) -> bool: + return any(eager_input_streaming_flag(tool) is True for tool in tools or ()) + @staticmethod def _supports_sampling_params(model: str) -> bool: """Claude 4.7+ (Opus 4.7/4.8, Fable 5) removed sampling params: the API diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 7eb56ae55d3..d1047c8d86c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -111,6 +111,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking_budget, ) from litellm.llms.anthropic.common_utils import ( + eager_input_streaming_flag, is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, strip_encrypted_reasoning_blocks_from_anthropic_messages, @@ -770,6 +771,7 @@ class LiteLLMAnthropicMessagesAdapter: "cache_control", "strict", "type", + "eager_input_streaming", ] for idx, tool in enumerate(tools): @@ -808,7 +810,14 @@ class LiteLLMAnthropicMessagesAdapter: for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs function_chunk.setdefault("parameters", {}).update({k: v}) - tool_param = ChatCompletionToolParam(type="function", function=function_chunk) + eager_input_streaming = eager_input_streaming_flag(tool) + tool_param = ( + ChatCompletionToolParam(type="function", function=function_chunk) + if eager_input_streaming is None + else ChatCompletionToolParam( + type="function", function=function_chunk, eager_input_streaming=eager_input_streaming + ) + ) self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 1a4f27e2ddd..87e28054817 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -48,6 +48,7 @@ from litellm.llms.bedrock.request_metadata import ( merge_bedrock_invoke_headers, resolve_bedrock_request_metadata, ) +from litellm.types.llms.anthropic import ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER from litellm.types.llms.bedrock import * from litellm.types.llms.openai import ( AllMessageValues, @@ -1518,10 +1519,7 @@ class AmazonConverseConfig(BaseConfig): bedrock_tools: list[ToolBlock] = [] # Collect anthropic_beta values from user headers - anthropic_beta_list: Final = [] - if headers: - user_betas: Final = get_anthropic_beta_from_headers(headers) - anthropic_beta_list.extend(user_betas) + anthropic_beta_list: Final = list(get_anthropic_beta_from_headers(headers or {})) # Separate pre-formatted Bedrock tools (e.g. systemTool from web_search_options) # from OpenAI-format tools that need transformation via _bedrock_tools_pt @@ -1634,6 +1632,12 @@ class AmazonConverseConfig(BaseConfig): if ANTHROPIC_EFFORT_BETA_HEADER not in anthropic_beta_list: anthropic_beta_list.append(ANTHROPIC_EFFORT_BETA_HEADER) + if ( + AnthropicModelInfo().is_eager_input_streaming_used(filtered_tools) + and ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER not in anthropic_beta_list + ): + anthropic_beta_list.append(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + # Bedrock Converse: compact_20260112 edits only (+ beta header). AmazonConverseConfig._filter_context_management_for_bedrock_converse( additional_request_params, anthropic_beta_list diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 72bc43ba938..1326dc22ca0 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -24,8 +24,12 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, strip_unsupported_bedrock_invoke_output_config_keys, + tools_without_eager_input_streaming, +) +from litellm.types.llms.anthropic import ( + ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, + ANTHROPIC_TOOL_SEARCH_BETA_HEADER, ) -from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -237,6 +241,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it) normalize_custom_field_on_tools(anthropic_request) normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_request) + outbound_tools: Final = tools_without_eager_input_streaming(anthropic_request) + if outbound_tools is not None: + anthropic_request["tools"] = outbound_tools return anthropic_request def _compute_bedrock_invoke_beta_headers( @@ -269,6 +276,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): if bedrock_supports_tool_search(model): beta_set.add("tool-search-tool-2025-10-19") + if self.is_eager_input_streaming_used(tools): + beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + auto_beta_list: Final = filter_and_transform_beta_headers( beta_headers=list(beta_set - user_beta_set), provider="bedrock", diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 62b485588ac..7e24292a87e 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -9,7 +9,7 @@ import functools import json import os import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict if TYPE_CHECKING: @@ -18,6 +18,7 @@ if TYPE_CHECKING: from litellm.types.llms.bedrock import BedrockCreateBatchRequest import httpx +from pydantic import TypeAdapter, ValidationError import litellm from litellm import verbose_logger @@ -330,6 +331,17 @@ def normalize_custom_field_on_tools(request_body: dict) -> None: tool["defer_loading"] = deferred +_TOOL_DICTS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def tools_without_eager_input_streaming(request_body: Mapping[str, object]) -> Sequence[object] | None: + try: + tools: Final = _TOOL_DICTS_ADAPTER.validate_python(request_body.get("tools")) + except ValidationError: + return None + return [{key: value for key, value in tool.items() if key != "eager_input_streaming"} for tool in tools] + + def normalize_json_schema_custom_types_to_object(schema: dict) -> None: """ In-place: replace JSON Schema ``type: \"custom\"`` with ``\"object\"`` (iterative walk). diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 4aa2afdbc78..d2be1ad9156 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -39,6 +39,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, strip_unsupported_bedrock_invoke_output_config_keys, + tools_without_eager_input_streaming, ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, @@ -46,6 +47,7 @@ from litellm.llms.bedrock.request_metadata import ( ) from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, + ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_BETA_HEADER, ) from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest @@ -525,6 +527,9 @@ class AmazonAnthropicClaudeMessagesConfig( if injected_thinking_for_clear_thinking: beta_set.add("interleaved-thinking-2025-05-14") + if anthropic_model_info.is_eager_input_streaming_used(tools): + beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + self._filter_context_management_for_bedrock_invoke( anthropic_messages_request=anthropic_messages_request, beta_set=beta_set, @@ -719,6 +724,10 @@ class AmazonAnthropicClaudeMessagesConfig( if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas + outbound_tools: Final = tools_without_eager_input_streaming(anthropic_messages_request) + if outbound_tools is not None: + anthropic_messages_request["tools"] = outbound_tools + remaining_output_config: Final = anthropic_messages_request.get("output_config") if ( litellm.drop_params is True diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5270df7b468..64daa695316 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19370,7 +19370,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { @@ -32899,6 +32899,10 @@ "cache_control": { "$ref": "#/components/schemas/ChatCompletionCachedContent" }, + "eager_input_streaming": { + "title": "Eager Input Streaming", + "type": "boolean" + }, "function": { "$ref": "#/components/schemas/ChatCompletionToolParamFunctionChunk" }, @@ -32928,6 +32932,10 @@ "title": "Description", "type": "string" }, + "eager_input_streaming": { + "title": "Eager Input Streaming", + "type": "boolean" + }, "name": { "title": "Name", "type": "string" diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index bcdee86360e..bcd24695f25 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -56,6 +56,7 @@ class AnthropicMessagesTool(TypedDict, total=False): defer_loading: bool allowed_callers: list[str] | None input_examples: list[dict[str, Any]] | None + eager_input_streaming: ReadOnly[bool] class AnthropicComputerTool(TypedDict, total=False): @@ -755,6 +756,8 @@ ANTHROPIC_TOOL_SEARCH_BETA_HEADER: Final = "advanced-tool-use-2025-11-20" # Effort beta header constant ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24" +ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER: Final = "fine-grained-tool-streaming-2025-05-14" + # OAuth constants ANTHROPIC_OAUTH_TOKEN_PREFIX: Final = "sk-ant-oat" ANTHROPIC_OAUTH_BETA_HEADER: Final = "oauth-2025-04-20" diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e3eac9b9205..abb76527e82 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -992,6 +992,7 @@ class ChatCompletionToolParamFunctionChunk(TypedDict, total=False): description: str parameters: dict strict: bool + eager_input_streaming: ReadOnly[bool] class OpenAIChatCompletionToolParam(TypedDict): @@ -1002,6 +1003,7 @@ class OpenAIChatCompletionToolParam(TypedDict): class ChatCompletionToolParam(OpenAIChatCompletionToolParam, total=False): cache_control: ChatCompletionCachedContent allowed_callers: list[str] + eager_input_streaming: ReadOnly[bool] class Function(TypedDict, total=False): diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 269c351f866..1359bf40dcd 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6370,3 +6370,74 @@ def test_response_format_tool_path_skips_forced_tool_choice_when_unsupported(loc assert "tools" in result assert "tool_choice" not in result + + +def _eager_chat_tool(**extra): + return { + "type": "function", + "function": { + "name": "write_file", + "description": "Write a file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, + }, + **extra, + } + + +@pytest.mark.parametrize("flag", [True, False]) +def test_eager_input_streaming_passed_through_from_tool_top_level(flag): + mapped_tool, _ = AnthropicConfig()._map_tool_helper(_eager_chat_tool(eager_input_streaming=flag)) + + assert mapped_tool == { + "name": "write_file", + "description": "Write a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, + "type": "custom", + "eager_input_streaming": flag, + } + + +def test_eager_input_streaming_passed_through_from_function(): + tool = _eager_chat_tool() + tool["function"]["eager_input_streaming"] = True + + mapped_tool, _ = AnthropicConfig()._map_tool_helper(tool) + + assert mapped_tool["eager_input_streaming"] is True + assert "eager_input_streaming" not in mapped_tool["input_schema"] + + +def test_eager_input_streaming_absent_stays_absent(): + mapped_tool, _ = AnthropicConfig()._map_tool_helper(_eager_chat_tool()) + + assert "eager_input_streaming" not in mapped_tool + + +def test_eager_input_streaming_rejects_non_boolean(): + with pytest.raises(litellm.BadRequestError, match="eager_input_streaming must be a boolean"): + AnthropicConfig()._map_tool_helper(_eager_chat_tool(eager_input_streaming="true")) + + +def test_eager_input_streaming_not_set_on_computer_use_tool(): + computer_tool = { + "type": "computer_20250124", + "function": {"name": "computer", "parameters": {"display_width_px": 1024, "display_height_px": 768}}, + "eager_input_streaming": True, + } + + mapped_tool, _ = AnthropicConfig()._map_tool_helper(computer_tool) + + assert mapped_tool["type"] == "computer_20250124" + assert "eager_input_streaming" not in mapped_tool + + +def test_eager_input_streaming_reaches_anthropic_request_tools(): + result = AnthropicConfig().map_openai_params( + non_default_params={"tools": [_eager_chat_tool(eager_input_streaming=True)], "stream": True}, + optional_params={}, + model="claude-sonnet-5", + drop_params=False, + ) + + assert result["tools"][0]["eager_input_streaming"] is True + assert result["tools"][0]["name"] == "write_file" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index e6782b70d3e..2318ceb937a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -5017,3 +5017,45 @@ def test_redacted_thinking_blocks_never_carry_cache_control(): replayed: Final = outbound["messages"][1]["content"][0] assert replayed["type"] == "redacted_thinking" assert "cache_control" not in replayed + + +EAGER_INPUT_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +@pytest.mark.parametrize("flag", [True, False]) +def test_translate_anthropic_tools_to_openai_carries_eager_input_streaming_onto_tool(flag): + """The per-tool flag lands on the OpenAI tool object, never inside the JSON schema Bedrock sends as inputSchema.""" + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA, "eager_input_streaming": flag}] + + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + assert new_tools[0]["eager_input_streaming"] is flag + assert new_tools[0]["function"]["parameters"] == EAGER_INPUT_SCHEMA + assert "eager_input_streaming" not in new_tools[0]["function"] + + +def test_translate_anthropic_tools_to_openai_omits_unset_eager_input_streaming(): + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA}] + + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + assert "eager_input_streaming" not in new_tools[0] + assert "eager_input_streaming" not in new_tools[0]["function"]["parameters"] + + +def test_eager_input_streaming_tool_reaches_bedrock_converse_as_beta(): + """An Anthropic Messages request routed to bedrock/converse/ turns the flag into the fine-grained streaming beta.""" + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA, "eager_input_streaming": True}] + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + data: Final = AmazonConverseConfig()._transform_request_helper( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + system_content_blocks=[], + optional_params={"tools": new_tools}, + messages=[{"role": "user", "content": "write a big file"}], + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == ["fine-grained-tool-streaming-2025-05-14"] + assert data["toolConfig"]["tools"][0]["toolSpec"]["inputSchema"]["json"] == EAGER_INPUT_SCHEMA diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index ec0bf6b842a..1395489334e 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -859,3 +859,56 @@ def test_bedrock_chat_invoke_tool_search_beta_follows_model_map( ) assert result.get("anthropic_beta") == expected_betas + + +FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +def _chat_invoke_request_with_tools(tools, headers=None): + config = AmazonAnthropicClaudeConfig() + model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" + optional_params = config.map_openai_params( + non_default_params={"max_tokens": 64, "stream": True, "tools": tools}, + optional_params={}, + model=model, + drop_params=False, + ) + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "write a big file"}], + optional_params=optional_params, + litellm_params={}, + headers=headers or {}, + ) + + +def _eager_openai_tool(name, **extra): + return {"type": "function", "function": {"name": name, "parameters": EAGER_TOOL_SCHEMA}, **extra} + + +def test_bedrock_chat_invoke_eager_input_streaming_tool_adds_beta_and_strips_key(): + result = _chat_invoke_request_with_tools( + [_eager_openai_tool("write_file", eager_input_streaming=True), _eager_openai_tool("read_file")] + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + assert [tool["name"] for tool in result["tools"]] == ["write_file", "read_file"] + assert all("eager_input_streaming" not in tool for tool in result["tools"]) + assert result["tools"][0]["input_schema"] == EAGER_TOOL_SCHEMA + + +def test_bedrock_chat_invoke_eager_input_streaming_false_strips_key_without_beta(): + result = _chat_invoke_request_with_tools([_eager_openai_tool("write_file", eager_input_streaming=False)]) + + assert "anthropic_beta" not in result + assert "eager_input_streaming" not in result["tools"][0] + + +def test_bedrock_chat_invoke_eager_input_streaming_beta_not_duplicated_with_client_header(): + result = _chat_invoke_request_with_tools( + [_eager_openai_tool("write_file", eager_input_streaming=True)], + headers={"anthropic-beta": FINE_GRAINED_TOOL_STREAMING_BETA}, + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 96c78c1cf75..650427f2ae3 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -7400,3 +7400,109 @@ def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it( ) assert result.choices[0].message.tool_calls is None assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000} + + +FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +def _eager_openai_tool(**extra): + return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA}, **extra} + + +def _eager_openai_function_tool(**extra): + return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA, **extra}} + + +def _eager_anthropic_tool(**extra): + return {"name": "write_file", "input_schema": EAGER_TOOL_SCHEMA, **extra} + + +def _converse_request(model, tools, headers=None): + return AmazonConverseConfig()._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params={"tools": tools}, + messages=[{"role": "user", "content": "write a big file"}], + headers=headers, + ) + + +@pytest.mark.parametrize( + "tool", + [ + _eager_openai_tool(eager_input_streaming=True), + _eager_openai_function_tool(eager_input_streaming=True), + _eager_anthropic_tool(eager_input_streaming=True), + ], + ids=["openai_top_level", "openai_under_function", "anthropic_shape"], +) +def test_eager_input_streaming_tool_adds_fine_grained_tool_streaming_beta(tool): + data = _converse_request("us.anthropic.claude-sonnet-4-5-20250929-v1:0", [tool]) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + tool_spec = data["toolConfig"]["tools"][0]["toolSpec"] + assert tool_spec["name"] == "write_file" + assert "eager_input_streaming" not in tool_spec + assert "eager_input_streaming" not in tool_spec["inputSchema"]["json"] + + +@pytest.mark.parametrize( + "tool", + [ + _eager_openai_tool(eager_input_streaming=False), + _eager_openai_function_tool(eager_input_streaming=False), + _eager_anthropic_tool(eager_input_streaming=False), + _eager_openai_tool(), + ], + ids=["openai_false", "function_false", "anthropic_false", "absent"], +) +def test_eager_input_streaming_false_or_absent_adds_no_beta(tool): + data = _converse_request("us.anthropic.claude-sonnet-4-5-20250929-v1:0", [tool]) + + assert "anthropic_beta" not in data.get("additionalModelRequestFields", {}) + assert "eager_input_streaming" not in data["toolConfig"]["tools"][0]["toolSpec"] + + +def test_eager_input_streaming_beta_only_on_anthropic_models(): + data = _converse_request("amazon.nova-pro-v1:0", [_eager_openai_tool(eager_input_streaming=True)]) + + assert "anthropic_beta" not in data.get("additionalModelRequestFields", {}) + assert data["toolConfig"]["tools"][0]["toolSpec"]["name"] == "write_file" + + +def test_eager_input_streaming_beta_not_duplicated_with_client_header(): + data = _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming=True)], + headers={"anthropic-beta": f"{FINE_GRAINED_TOOL_STREAMING_BETA},interleaved-thinking-2025-05-14"}, + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [ + FINE_GRAINED_TOOL_STREAMING_BETA, + "interleaved-thinking-2025-05-14", + ] + + +def test_eager_input_streaming_beta_never_written_back_into_client_header_list(): + headers = {"anthropic-beta": ["interleaved-thinking-2025-05-14"]} + + data = _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming=True)], + headers=headers, + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [ + "interleaved-thinking-2025-05-14", + FINE_GRAINED_TOOL_STREAMING_BETA, + ] + assert headers == {"anthropic-beta": ["interleaved-thinking-2025-05-14"]} + + +def test_eager_input_streaming_non_boolean_is_a_bad_request(): + with pytest.raises(litellm.BadRequestError, match="eager_input_streaming must be a boolean"): + _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming="true")], + ) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index e43accdb835..b71e8e521f2 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3244,3 +3244,65 @@ def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_mo ) assert result.get("output_config") == {"format": schema_format} + + +FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" + + +def _invoke_request_with_tools(tools, headers=None): + from litellm.types.router import GenericLiteLLMParams + + return AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "write a big file"}], + anthropic_messages_optional_request_params={"max_tokens": 4096, "tools": copy.deepcopy(tools), "stream": True}, + litellm_params=GenericLiteLLMParams(), + headers=headers or {}, + ) + + +def _eager_invoke_tool(name, eager_input_streaming): + return { + "name": name, + "description": f"{name} tool", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + "eager_input_streaming": eager_input_streaming, + } + + +def test_bedrock_invoke_eager_input_streaming_tool_adds_beta_and_strips_key(): + result = _invoke_request_with_tools( + [ + _eager_invoke_tool("write_file", True), + _eager_invoke_tool("read_file", False), + {"name": "list_files", "input_schema": {"type": "object", "properties": {}}}, + ] + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + assert [tool["name"] for tool in result["tools"]] == ["write_file", "read_file", "list_files"] + assert all("eager_input_streaming" not in tool for tool in result["tools"]) + assert result["tools"][0]["description"] == "write_file tool" + assert result["tools"][0]["input_schema"] == {"type": "object", "properties": {"path": {"type": "string"}}} + + +def test_bedrock_invoke_eager_input_streaming_false_strips_key_without_beta(): + result = _invoke_request_with_tools([_eager_invoke_tool("write_file", False)]) + + assert "anthropic_beta" not in result + assert result["tools"] == [ + { + "name": "write_file", + "description": "write_file tool", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + } + ] + + +def test_bedrock_invoke_eager_input_streaming_beta_not_duplicated_with_client_header(): + result = _invoke_request_with_tools( + [_eager_invoke_tool("write_file", True)], + headers={"anthropic-beta": FINE_GRAINED_TOOL_STREAMING_BETA}, + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1b0081cb9c..cc5169dd84c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25970,6 +25970,8 @@ export interface components { /** Allowed Callers */ allowed_callers?: string[]; cache_control?: components["schemas"]["ChatCompletionCachedContent"]; + /** Eager Input Streaming */ + eager_input_streaming?: boolean; function: components["schemas"]["ChatCompletionToolParamFunctionChunk"]; /** Type */ type: "function" | string; @@ -25978,6 +25980,8 @@ export interface components { ChatCompletionToolParamFunctionChunk: { /** Description */ description?: string; + /** Eager Input Streaming */ + eager_input_streaming?: boolean; /** Name */ name: string; /** Parameters */ From 817aefc41348aa9f47b246b554afc5efbf6e2523 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 20:15:48 +0000 Subject: [PATCH 059/102] test(proxy): keep stored pass-through route tests off the shared FastAPI app so the allowlist coverage test stays order independent Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_server.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index bc7ae556e7b..5246463f616 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7484,7 +7484,8 @@ async def test_deleting_the_stored_pass_through_row_takes_the_route_out_of_servi settings: Final = patch("litellm.proxy.proxy_server.general_settings", {}) # test-quality-ok: the method reads this module global; no injection seam yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", None) # test-quality-ok: module global holding the YAML endpoints; this case has none - with settings, yaml_endpoints: + app_routes: Final = patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.SafeRouteAdder.add_api_route_if_not_exists") # test-quality-ok: the registry is the observable; a real route would stay on the shared FastAPI app for the rest of the xdist worker + with settings, yaml_endpoints, app_routes: pc = ProxyConfig() await pc._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]}) assert live_routes(), "the stored endpoint should be serving before the row is deleted" @@ -7517,7 +7518,8 @@ async def test_a_stored_pass_through_row_never_disturbs_the_config_declared_rout settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [config_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [config_endpoint]) # test-quality-ok: module global holding the YAML endpoints the reload merges in - with settings, yaml_endpoints: + app_routes: Final = patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.SafeRouteAdder.add_api_route_if_not_exists") # test-quality-ok: the registry is the observable; a real route would stay on the shared FastAPI app for the rest of the xdist worker + with settings, yaml_endpoints, app_routes: await initialize_pass_through_endpoints(pass_through_endpoints=[config_endpoint]) assert live_paths() == {config_path} From c54c0b049daf75fd761723d388e81b9b9cbb8f9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:16:56 -0700 Subject: [PATCH 060/102] fix(vertex_ai): keep reasoning_effort unsupported on Vertex AI Mistral partner models Vertex AI Mistral models reused MistralConfig, whose reasoning_effort advertisement checks the mistral provider entry of the cost map, so vertex_ai/mistral-medium-3 started advertising reasoning_effort and drop_params stopped dropping it, turning a 200 into a Vertex 400. VertexAIMistralConfig scopes that lookup to the vertex_ai provider, and MistralConfig now reads the provider from its custom_llm_provider property instead of a hardcoded "mistral". --- litellm/__init__.py | 3 +++ litellm/_lazy_imports_registry.py | 5 +++++ .../get_supported_openai_params.py | 2 +- litellm/llms/mistral/chat/transformation.py | 21 ++++++++++++++----- .../mistral/transformation.py | 7 +++++++ litellm/utils.py | 4 ++-- ...i_partner_models_mistral_transformation.py | 17 +++++++++++++++ 7 files changed, 51 insertions(+), 8 deletions(-) create mode 100644 litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/transformation.py create mode 100644 tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 71857877e53..2e7a2a4efd8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1704,6 +1704,9 @@ if TYPE_CHECKING: from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import ( VertexAIAi21Config as VertexAIAi21Config, ) + from .llms.vertex_ai.vertex_ai_partner_models.mistral.transformation import ( + VertexAIMistralConfig as VertexAIMistralConfig, + ) from .llms.bedrock.chat.invoke_handler import ( AmazonCohereChatConfig as AmazonCohereChatConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 4c478b51ed1..9cfcb9e41f7 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -184,6 +184,7 @@ LLM_CONFIG_NAMES: Final = ( "VertexAIAnthropicConfig", "VertexAILlama3Config", "VertexAIAi21Config", + "VertexAIMistralConfig", "AmazonCohereChatConfig", "AmazonBedrockGlobalConfig", "AmazonAI21Config", @@ -771,6 +772,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.vertex_ai.vertex_ai_partner_models.ai21.transformation", "VertexAIAi21Config", ), + "VertexAIMistralConfig": ( + ".llms.vertex_ai.vertex_ai_partner_models.mistral.transformation", + "VertexAIMistralConfig", + ), "AmazonCohereChatConfig": ( ".llms.bedrock.chat.invoke_handler", "AmazonCohereChatConfig", diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 915a03025d9..08b8816e17d 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -190,7 +190,7 @@ def get_supported_openai_params( elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta": if request_type == "chat_completion": if model.startswith("mistral"): - return litellm.MistralConfig().get_supported_openai_params(model=model) + return litellm.VertexAIMistralConfig().get_supported_openai_params(model=model) elif model.startswith("codestral"): return litellm.CodestralTextCompletionConfig().get_supported_openai_params(model=model) elif model.startswith("claude"): diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 8cbe4291054..f77e828b59a 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -35,14 +35,19 @@ if TYPE_CHECKING: import tiktoken -def _accepted_reasoning_effort(model: str, requested: str) -> str: - declared: Final = declared_reasoning_efforts_for_model(model, "mistral") +def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider: str) -> str: + declared: Final = declared_reasoning_efforts_for_model(model, custom_llm_provider) if declared is None: return requested accepted: Final = nearest_declared_reasoning_effort(requested, declared) if accepted != requested: verbose_logger.debug( - "mistral: %s takes reasoning_effort %s, sending %s in place of %s", model, declared, accepted, requested + "%s: %s takes reasoning_effort %s, sending %s in place of %s", + custom_llm_provider, + model, + declared, + accepted, + requested, ) return accepted @@ -103,9 +108,15 @@ class MistralConfig(OpenAIGPTConfig): def get_config(cls): return super().get_config() + @property + def custom_llm_provider(self) -> str: + return "mistral" + def get_supported_openai_params(self, model: str) -> list[str]: is_magistral: Final = "magistral" in model.lower() - accepts_reasoning_effort: Final = is_magistral or supports_reasoning(model=model, custom_llm_provider="mistral") + accepts_reasoning_effort: Final = is_magistral or supports_reasoning( + model=model, custom_llm_provider=self.custom_llm_provider + ) return [ "stream", "temperature", @@ -187,7 +198,7 @@ class MistralConfig(OpenAIGPTConfig): if param == "response_format": optional_params["response_format"] = value if param == "reasoning_effort" and "magistral" not in model.lower(): - optional_params["reasoning_effort"] = _accepted_reasoning_effort(model, value) + optional_params["reasoning_effort"] = _accepted_reasoning_effort(model, value, self.custom_llm_provider) if param in ("reasoning_effort", "thinking") and "magistral" in model.lower(): # Flag that we need to add reasoning system prompt optional_params["_add_reasoning_prompt"] = True diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/transformation.py new file mode 100644 index 00000000000..18321c8768a --- /dev/null +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/transformation.py @@ -0,0 +1,7 @@ +from litellm.llms.mistral.chat.transformation import MistralConfig + + +class VertexAIMistralConfig(MistralConfig): + @property + def custom_llm_provider(self) -> str: + return "vertex_ai" diff --git a/litellm/utils.py b/litellm/utils.py index 2c9200fbad7..56985a11b8c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4510,7 +4510,7 @@ def get_optional_params( drop_params=bool(drop_params), ) else: - optional_params = litellm.MistralConfig().map_openai_params( + optional_params = litellm.VertexAIMistralConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, @@ -8380,7 +8380,7 @@ class ProviderConfigManager: elif model in litellm.vertex_mistral_models: if "codestral" in model: return litellm.CodestralTextCompletionConfig() - return litellm.MistralConfig() + return litellm.VertexAIMistralConfig() elif model in litellm.vertex_ai_ai21_models: return litellm.VertexAIAi21Config() else: diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py new file mode 100644 index 00000000000..f7df4507651 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/mistral/test_vertex_ai_partner_models_mistral_transformation.py @@ -0,0 +1,17 @@ +import litellm + + +def test_reasoning_effort_stays_unsupported_on_vertex_partner_models(local_model_cost_map): + assert "reasoning_effort" in litellm.get_supported_openai_params( + model="mistral-medium-3", custom_llm_provider="mistral" + ) + assert "reasoning_effort" not in litellm.get_supported_openai_params( + model="mistral-medium-3", custom_llm_provider="vertex_ai" + ) + dropped = litellm.get_optional_params( + model="mistral-medium-3", + custom_llm_provider="vertex_ai", + reasoning_effort="high", + drop_params=True, + ) + assert "reasoning_effort" not in dropped From 0f59e61f24dc72c6f4a9af15f1618d1b13591cde Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:17:11 +0000 Subject: [PATCH 061/102] fix(rust): align exception mapping tests with Python Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/exception_mapping_utils/mod.rs | 25 +++++++++---------- .../src/exception_mapping_utils/rules.rs | 4 +-- 2 files changed, 14 insertions(+), 15 deletions(-) diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs index 8892fba1399..acc7820c770 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs @@ -282,7 +282,6 @@ fn extra_information(context: &ExceptionContext, api_base: Option<&str>) -> Stri let lines = [ Some(format!("\nModel: {}", context.model)), api_base.map(|api_base| format!("\nAPI Base: `{api_base}`")), - (!context.redact_messages_in_exceptions).then(|| "\nMessages: `None`".to_string()), context .model_group .as_ref() @@ -314,7 +313,7 @@ fn extra_information(context: &ExceptionContext, api_base: Option<&str>) -> Stri mod testing { use super::*; - pub(super) const DEBUG: &str = "\nModel: ocr-model\nMessages: `None`"; + pub(super) const DEBUG: &str = "\nModel: ocr-model"; pub(super) fn context(provider: &str, family: ExceptionFamily) -> ExceptionContext { ExceptionContext { @@ -369,14 +368,6 @@ mod testing { ..failure } } - - /// What `exception_type` returns for an `http` original the rule mapped to `failure`. - pub(super) fn with_headers(failure: PublicFailure) -> PublicFailure { - PublicFailure { - litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), - ..failure - } - } } #[cfg(test)] @@ -492,7 +483,17 @@ mod tests { #[case] message: &str, ) { let context = context("reducto", ExceptionFamily::Other); - let expected = failure(PublicKind::ApiConnection, message, "reducto"); + let expected = match original { + OriginalException::Http { .. } => with_debug(failure( + status(StatusClass::BadRequest, upstream(409, "rejected")), + message, + "reducto", + )), + OriginalException::Connection { .. } | OriginalException::Local { .. } => { + failure(PublicKind::ApiConnection, message, "reducto") + } + _ => unreachable!(), + }; let actual = exception_type(&context, &original); assert_eq!( PublicFailure { @@ -669,7 +670,6 @@ mod tests { "\n\nKey Name: `key`\nTeam: `None`", "\nModel: ocr-model", "\nAPI Base: `region-aiplatform.googleapis.com/v1/projects/project/locations/region/publishers/google/models/ocr-model:generateContent`", - "\nMessages: `None`", "\nmodel_group: `ocr`\n", "\ndeployment: `deployment`\n", "\nvertex_project: `project`\n", @@ -681,7 +681,6 @@ mod tests { #[rstest::rstest] #[case::bare(ExceptionContext::default(), "\nModel: ")] #[case::redacted_messages(ExceptionContext { redact_messages_in_exceptions: true, model: "m".into(), ..ExceptionContext::default() }, "\nModel: m")] - #[case::messages(ExceptionContext { model: "m".into(), ..ExceptionContext::default() }, "\nModel: m\nMessages: `None`")] #[case::team_alias( ExceptionContext { model: "m".into(), user_api_key_alias: Some("key".into()), user_api_key_team_alias: Some("team".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, "\n\nKey Name: `key`\nTeam: `team`\nModel: m" diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs index 5df8ad991b5..34cf2c390c3 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs @@ -224,7 +224,7 @@ mod tests { "mistral", ))] #[case::later_rule_when_the_earlier_does_not_apply("second", PublicFailure { - litellm_debug_info: Some("\nModel: ocr-model\nMessages: `None`".into()), + litellm_debug_info: Some("\nModel: ocr-model".into()), ..failure(PublicKind::ApiConnection, "seen second", "mistral") })] fn the_first_applicable_rule_decides(#[case] body: &str, #[case] expected: PublicFailure) { @@ -324,7 +324,7 @@ mod tests { } #[rstest::rstest] - #[case::with_debug(true, Some("\nModel: ocr-model\nMessages: `None`"))] + #[case::with_debug(true, Some("\nModel: ocr-model"))] #[case::without_debug(false, None)] fn debug_rules_carry_the_extra_information( #[case] debug: bool, From 2703fab3c810f1c82b65c035acf4920ad3ac85e3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:21:21 +0000 Subject: [PATCH 062/102] fix(deps): update anyio for osv scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- uv.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index a5e60c68515..29027c2a75f 100644 --- a/uv.lock +++ b/uv.lock @@ -315,16 +315,16 @@ vertex = [ [[package]] name = "anyio" -version = "4.13.0" +version = "4.15.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, { name = "idna" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/d2/f4d173e22df740bc37b1db102b386ba719b66e95b0f0d751f556b387e6d2/anyio-4.15.1.tar.gz", hash = "sha256:9f28306018cbd6d329e64a36d58256edff76dd996fe423bc957326e578b82a94", size = 276966, upload-time = "2026-09-05T10:42:39.44Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, + { url = "https://files.pythonhosted.org/packages/12/b8/4bd346e22b28902df4d651910f5242c28d84e4a5c2435ca5c3f797ed7e2e/anyio-4.15.1-py3-none-any.whl", hash = "sha256:6152fdbbf9a77fdec97731721bebf7c4c44f7c29b424b0065826173efc7ed101", size = 132079, upload-time = "2026-09-05T10:42:37.923Z" }, ] [[package]] From e036e256ef00be1afbed0cfb0a9ce87f67d18fa5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:23:53 +0000 Subject: [PATCH 063/102] revert: drop redundant anyio lockfile bump Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- uv.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index 29027c2a75f..a5e60c68515 100644 --- a/uv.lock +++ b/uv.lock @@ -315,16 +315,16 @@ vertex = [ [[package]] name = "anyio" -version = "4.15.1" +version = "4.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, { name = "idna" }, - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a9/d2/f4d173e22df740bc37b1db102b386ba719b66e95b0f0d751f556b387e6d2/anyio-4.15.1.tar.gz", hash = "sha256:9f28306018cbd6d329e64a36d58256edff76dd996fe423bc957326e578b82a94", size = 276966, upload-time = "2026-09-05T10:42:39.44Z" } +sdist = { url = "https://files.pythonhosted.org/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/12/b8/4bd346e22b28902df4d651910f5242c28d84e4a5c2435ca5c3f797ed7e2e/anyio-4.15.1-py3-none-any.whl", hash = "sha256:6152fdbbf9a77fdec97731721bebf7c4c44f7c29b424b0065826173efc7ed101", size = 132079, upload-time = "2026-09-05T10:42:37.923Z" }, + { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, ] [[package]] From a37d1e8b9b752f2c0672bc394e14095ef7c69b31 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:30:08 -0700 Subject: [PATCH 064/102] test: type the eager_input_streaming test helpers and mark constants Final --- .../test_anthropic_chat_transformation.py | 22 +++++++++---------- ...ations_anthropic_claude3_transformation.py | 17 ++++++++------ .../chat/test_converse_transformation.py | 15 ++++++++----- .../test_anthropic_claude3_transformation.py | 9 +++++--- 4 files changed, 36 insertions(+), 27 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 1359bf40dcd..5e179f950a0 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6372,18 +6372,19 @@ def test_response_format_tool_path_skips_forced_tool_choice_when_unsupported(loc assert "tool_choice" not in result -def _eager_chat_tool(**extra): +def _eager_chat_function(**extra: object) -> dict[str, object]: return { - "type": "function", - "function": { - "name": "write_file", - "description": "Write a file", - "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, - }, + "name": "write_file", + "description": "Write a file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, **extra, } +def _eager_chat_tool(**extra: object) -> dict[str, object]: + return {"type": "function", "function": _eager_chat_function(), **extra} + + @pytest.mark.parametrize("flag", [True, False]) def test_eager_input_streaming_passed_through_from_tool_top_level(flag): mapped_tool, _ = AnthropicConfig()._map_tool_helper(_eager_chat_tool(eager_input_streaming=flag)) @@ -6398,10 +6399,9 @@ def test_eager_input_streaming_passed_through_from_tool_top_level(flag): def test_eager_input_streaming_passed_through_from_function(): - tool = _eager_chat_tool() - tool["function"]["eager_input_streaming"] = True - - mapped_tool, _ = AnthropicConfig()._map_tool_helper(tool) + mapped_tool, _ = AnthropicConfig()._map_tool_helper( + {"type": "function", "function": _eager_chat_function(eager_input_streaming=True)} + ) assert mapped_tool["eager_input_streaming"] is True assert "eager_input_streaming" not in mapped_tool["input_schema"] diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 1395489334e..84db0733227 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,6 +1,7 @@ import asyncio import json import uuid +from typing import Final from unittest.mock import patch import httpx @@ -861,14 +862,16 @@ def test_bedrock_chat_invoke_tool_search_beta_follows_model_map( assert result.get("anthropic_beta") == expected_betas -FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" -EAGER_TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} -def _chat_invoke_request_with_tools(tools, headers=None): - config = AmazonAnthropicClaudeConfig() - model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" - optional_params = config.map_openai_params( +def _chat_invoke_request_with_tools( + tools: list[dict[str, object]], headers: dict[str, str] | None = None +) -> dict[str, object]: + config: Final = AmazonAnthropicClaudeConfig() + model: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" + optional_params: Final = config.map_openai_params( non_default_params={"max_tokens": 64, "stream": True, "tools": tools}, optional_params={}, model=model, @@ -883,7 +886,7 @@ def _chat_invoke_request_with_tools(tools, headers=None): ) -def _eager_openai_tool(name, **extra): +def _eager_openai_tool(name: str, **extra: object) -> dict[str, object]: return {"type": "function", "function": {"name": name, "parameters": EAGER_TOOL_SCHEMA}, **extra} diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 650427f2ae3..cd33e8d34d8 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -4,6 +4,7 @@ import os import httpx import pytest +from typing import Final from unittest.mock import MagicMock, patch import litellm @@ -7402,23 +7403,25 @@ def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it( assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000} -FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" -EAGER_TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} -def _eager_openai_tool(**extra): +def _eager_openai_tool(**extra: object) -> dict[str, object]: return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA}, **extra} -def _eager_openai_function_tool(**extra): +def _eager_openai_function_tool(**extra: object) -> dict[str, object]: return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA, **extra}} -def _eager_anthropic_tool(**extra): +def _eager_anthropic_tool(**extra: object) -> dict[str, object]: return {"name": "write_file", "input_schema": EAGER_TOOL_SCHEMA, **extra} -def _converse_request(model, tools, headers=None): +def _converse_request( + model: str, tools: list[dict[str, object]], headers: dict[str, object] | None = None +) -> dict[str, object]: return AmazonConverseConfig()._transform_request_helper( model=model, system_content_blocks=[], diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index b71e8e521f2..7be005c0efe 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -4,6 +4,7 @@ import json import os from datetime import datetime from types import SimpleNamespace +from typing import Final from unittest.mock import Mock import pytest @@ -3246,10 +3247,12 @@ def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_mo assert result.get("output_config") == {"format": schema_format} -FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14" +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" -def _invoke_request_with_tools(tools, headers=None): +def _invoke_request_with_tools( + tools: list[dict[str, object]], headers: dict[str, str] | None = None +) -> dict[str, object]: from litellm.types.router import GenericLiteLLMParams return AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( @@ -3261,7 +3264,7 @@ def _invoke_request_with_tools(tools, headers=None): ) -def _eager_invoke_tool(name, eager_input_streaming): +def _eager_invoke_tool(name: str, eager_input_streaming: bool) -> dict[str, object]: return { "name": name, "description": f"{name} tool", From e5fd2bec832be2ddf76851f73859d8bc2fb2b957 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:34:29 -0700 Subject: [PATCH 065/102] fix: forward eager_input_streaming through the Responses API tool bridge --- .../transformation.py | 4 +++ .../test_litellm_completion_responses.py | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 01fb6cb483d..0d1297f63e2 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2020,6 +2020,8 @@ class LiteLLMCompletionResponsesConfig: chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") if tool.get("input_examples"): chat_completion_tool["input_examples"] = tool.get("input_examples") + if tool.get("eager_input_streaming") is not None: + chat_completion_tool["eager_input_streaming"] = tool.get("eager_input_streaming") return ResponsesToolChatForm( chat_tools=(cast(ChatCompletionToolParam, chat_completion_tool),), web_search_options=None ) @@ -2096,6 +2098,8 @@ class LiteLLMCompletionResponsesConfig: responses_tool["allowed_callers"] = tool.get("allowed_callers") if tool.get("input_examples") is not None: responses_tool["input_examples"] = tool.get("input_examples") + if tool.get("eager_input_streaming") is not None: + responses_tool["eager_input_streaming"] = tool.get("eager_input_streaming") result.append(responses_tool) else: # mcp or other: pass through unchanged diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..f1924b87a20 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1900,6 +1900,38 @@ class TestToolTransformation: assert "defer_loading" not in result_tool assert "allowed_callers" not in result_tool assert "input_examples" not in result_tool + assert "eager_input_streaming" not in result_tool + + @pytest.mark.parametrize("eager_input_streaming", [True, False]) + def test_transform_function_tools_forwards_eager_input_streaming(self, eager_input_streaming: bool) -> None: + function_tool: Final = { + "type": "function", + "name": "write_file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}, + "eager_input_streaming": eager_input_streaming, + } + + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[function_tool] + ) + + assert result_tools[0]["eager_input_streaming"] is eager_input_streaming + + @pytest.mark.parametrize("eager_input_streaming", [True, False]) + def test_chat_completion_tools_to_responses_tools_keeps_eager_input_streaming( + self, eager_input_streaming: bool + ) -> None: + chat_tool: Final = { + "type": "function", + "function": {"name": "write_file", "parameters": {"type": "object"}}, + "eager_input_streaming": eager_input_streaming, + } + + result_tools: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( + [chat_tool] + ) + + assert result_tools[0]["eager_input_streaming"] is eager_input_streaming def test_transform_code_execution_tools(self): """Test that code_execution tools are passed through as-is""" From abf530fbeb0563d5877ceb883261b25161855b29 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 20:55:11 +0000 Subject: [PATCH 066/102] fix(proxy): never rewind the daily global spend marker from an overlapping reconcile run Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../daily_global_spend_rollup.py | 31 +++++++++++++------ .../test_daily_global_spend_rollup.py | 30 +++++++++++++++++- 2 files changed, 51 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index c135c7d1d9c..d50feb6f16b 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -161,6 +161,24 @@ async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThroug await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) +async def _stored_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: + """The marker as another pod may have just written it, bypassing this pod's config cache.""" + param: Final = await ConfigRepository(prisma_client).get_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + return None if param is None else _marker_from_param_value(param.param_value) + + +async def _record_advanced(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: + """Advance the stored marker by ``days``. Two runs can overlap (Redis unreachable, lock expired + on a long backfill), so the base is what is stored now, not the snapshot this run scanned from: + a slower run may then only add to the faster run's marker, never rewind it. Without a new scan + time the stored one is kept.""" + stored: Final = await _stored_marker(prisma_client) + kept_scanned_at: Final = None if stored is None else stored.scanned_at + await _record_marker( + prisma_client, _advanced(stored, days, scanned_at=scanned_at if scanned_at is not None else kept_scanned_at) + ) + + async def _db_now(prisma_client: "PrismaClient") -> _NowRow: rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) return _NowRow.model_validate(rows[0]) @@ -203,7 +221,7 @@ async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> Rec marker: Final = await reconciled_through(prisma_client) return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) if scan.marker is not None or done: - await _record_marker(prisma_client, _advanced(scan.marker, done, scanned_at=scan.scanned_at)) + await _record_advanced(prisma_client, done, scanned_at=scan.scanned_at) return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) @@ -215,21 +233,16 @@ def _advanced(marker: ReconciledThrough | None, days: tuple[str, ...], *, scanne async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: for index, day in enumerate(scan.days): - if not await _reconcile_and_record(prisma_client, scan.marker, scan.days[: index + 1]): + if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]): return scan.days[:index] return scan.days -async def _reconcile_and_record( - prisma_client: "PrismaClient", marker: ReconciledThrough | None, done_with_this: tuple[str, ...] -) -> bool: +async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: tuple[str, ...]) -> bool: day: Final = done_with_this[-1] try: await reconcile_day(prisma_client, day) - await _record_marker( - prisma_client, - _advanced(marker, done_with_this, scanned_at=None if marker is None else marker.scanned_at), - ) + await _record_advanced(prisma_client, done_with_this, scanned_at=None) except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) return False diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 9655953134a..69f06fad081 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -40,6 +40,10 @@ class _FakeConfigTable: self.rows[where["param_name"]] = data["update"]["param_value"] return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) + async def find_unique(self, *, where: dict[str, str]) -> _FakeConfigRow | None: + stored = self.rows.get(where["param_name"]) + return None if stored is None else _FakeConfigRow(where["param_name"], stored) + class _FakeDb: """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, @@ -68,11 +72,15 @@ class _FakeDb: if day in self._prisma.failing_days: raise RuntimeError(f"day {day} exploded") self._prisma.reconciled.append(day) + landing = self._prisma.marker_landing_on_day.get(day) + if landing is not None: + self.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = landing return 1 class _FakePrisma: - """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" + """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw. + ``marker_landing_on_day`` stores another pod's marker the moment this run rewrites that day.""" def __init__( self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset(), today: date = TODAY @@ -81,6 +89,7 @@ class _FakePrisma: self.today = today self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} self.failing_days = failing_days + self.marker_landing_on_day: dict[str, str] = {} self.reconciled: list[str] = [] self.db = _FakeDb(self) @@ -221,6 +230,25 @@ async def test_the_next_run_resumes_from_the_failed_day(): assert await reconciled_through(prisma) == "2026-09-03" +@pytest.mark.asyncio +async def test_a_slower_overlapping_run_never_rewinds_the_marker_a_faster_run_stored(): + """Two pods can reconcile at once (Redis unreachable, or the lock expired on a long backfill). + When the faster one has already stored a later marker, the slower one may only add to it. Putting + its own older prefix back, or dropping the scan time, would send usage reads for every day in + between back to the per-key table until the next run.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-03"})) + prisma.marker_landing_on_day = { + "2026-09-02": '{"reconciled_through": "2026-09-14", "scanned_at": "clock-0009"}', + } + + result = await run_daily_global_spend_reconcile(prisma) + + assert result.days_reconciled == ("2026-09-01", "2026-09-02") + assert result.reconciled_through == "2026-09-14" + marker = await read_marker(prisma) + assert marker is not None and (marker.reconciled_through, marker.scanned_at) == ("2026-09-14", "clock-0009") + + @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" From 836bf7d8974bc3fc3cdd71a3fe6d9565201c2f2e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 13:59:36 -0700 Subject: [PATCH 067/102] refactor(rust): make the exception mapper a pure rule table Rebuild exception_type around text rules per provider family and one shared status table. The mapper takes the context, an injected redactor and the original failure, and returns a PublicError with the message, the real upstream response and the debug text. Every divergence from the Python mapper and every known gap is listed in the module header Match Python on a standalone 429 with an unknown status and on Cohere's rules for failures without a status. Drop python_repr and the unread public_failures fixtures --- litellm-rust/Cargo.lock | 1 - litellm-rust/crates/core-utils/Cargo.toml | 1 - .../src/exception_mapping_utils/cohere.rs | 275 ++----- .../src/exception_mapping_utils/mod.rs | 757 +++++++----------- .../src/exception_mapping_utils/openai.rs | 600 +++----------- .../src/exception_mapping_utils/original.rs | 145 +++- .../src/exception_mapping_utils/public.rs | 257 +----- .../src/exception_mapping_utils/rules.rs | 283 +------ .../src/exception_mapping_utils/status.rs | 182 +---- .../src/exception_mapping_utils/vertex_ai.rs | 573 ++----------- litellm-rust/crates/core-utils/src/lib.rs | 1 - .../crates/core-utils/src/python_repr.rs | 93 --- .../crates/core-utils/src/secret_redaction.rs | 50 +- .../fixtures/public_failures/api.json | 13 - .../public_failures/api_connection.json | 11 - .../status_authentication.json | 13 - .../public_failures/status_bad_gateway.json | 28 - .../public_failures/status_bad_request.json | 28 - .../status_content_policy_violation.json | 28 - .../status_context_window_exceeded.json | 28 - .../status_internal_server.json | 19 - .../public_failures/status_not_found.json | 28 - .../status_permission_denied.json | 19 - .../public_failures/status_rate_limit.json | 28 - .../status_service_unavailable.json | 28 - .../status_unsupported_params.json | 28 - .../public_failures/timeout_with_status.json | 12 - .../timeout_without_status.json | 12 - 28 files changed, 807 insertions(+), 2734 deletions(-) delete mode 100644 litellm-rust/crates/core-utils/src/python_repr.rs delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/api.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/api_connection.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_authentication.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_gateway.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_request.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_content_policy_violation.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_context_window_exceeded.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_internal_server.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_not_found.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_permission_denied.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_rate_limit.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_service_unavailable.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/status_unsupported_params.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_with_status.json delete mode 100644 tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_without_status.json diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 6c524124550..c359ca19986 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2092,7 +2092,6 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_with", - "strum", "thiserror 2.0.19", "url", ] diff --git a/litellm-rust/crates/core-utils/Cargo.toml b/litellm-rust/crates/core-utils/Cargo.toml index 59e0ee1a09d..eb353bc060c 100644 --- a/litellm-rust/crates/core-utils/Cargo.toml +++ b/litellm-rust/crates/core-utils/Cargo.toml @@ -12,7 +12,6 @@ serde.workspace = true serde_json.workspace = true serde_path_to_error = "0.1" serde_with.workspace = true -strum.workspace = true thiserror.workspace = true url.workspace = true diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/cohere.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/cohere.rs index 381ebd7ad27..c2a391ee223 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/cohere.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/cohere.rs @@ -1,232 +1,115 @@ -use super::Mapping; -use super::public::{PublicFailure, StatusClass}; -use super::rules::{Kind, ResponseChoice, Rule, apply, contains_any}; +use super::public::PublicError; +use super::rules::{Rule, contains_any}; -const fn with_response(class: StatusClass) -> Kind { - Kind::Status { - class, - response: ResponseChoice::Provider, - } -} - -fn original(mapping: &Mapping<'_>) -> String { - format!("CohereException - {}", mapping.original.message) -} - -fn status_is(mapping: &Mapping<'_>, statuses: &[u16]) -> bool { - mapping - .original - .status - .is_some_and(|status| statuses.contains(&status)) -} - -/// `_map_cohere_exception`, in its branch order. A failure no rule claims falls through to -/// the status table. -const RULES: &[Rule] = &[ - Rule { - when: |mapping| { +/// The text branches of `_map_cohere_exception`, in its order. +pub(super) const RULES: &[Rule] = &[ + Rule::new( + |mapping| { contains_any( &mapping.error_str, &["invalid api token", "No API key provided."], ) }, - kind: with_response(StatusClass::Authentication), - message: original, - debug: false, - }, - Rule { - when: |mapping| mapping.error_str.contains("invalid type: parameter"), - kind: with_response(StatusClass::BadRequest), - message: original, - debug: false, - }, - Rule { - when: |mapping| mapping.error_str.contains("too many tokens"), - kind: with_response(StatusClass::ContextWindowExceeded), - message: original, - debug: false, - }, - Rule { - when: |mapping| { + PublicError::Authentication, + ), + Rule::new( + |mapping| mapping.error_str.contains("invalid type: parameter"), + PublicError::BadRequest, + ), + Rule::new( + |mapping| mapping.error_str.contains("too many tokens"), + PublicError::ContextWindowExceeded, + ), + Rule::new( + |mapping| { mapping .error_str .to_lowercase() .contains("internal server error") }, - kind: with_response(StatusClass::InternalServer), - message: |mapping| format!("CohereException - {}", mapping.error_str), - debug: false, - }, - Rule { - when: |mapping| status_is(mapping, &[400, 498]), - kind: with_response(StatusClass::BadRequest), - message: original, - debug: false, - }, - Rule { - when: |mapping| status_is(mapping, &[408]), - kind: Kind::Timeout(None), - message: original, - debug: false, - }, - Rule { - when: |mapping| status_is(mapping, &[500]), - kind: with_response(StatusClass::InternalServer), - message: original, - debug: false, - }, + PublicError::InternalServer, + ), + Rule::new( + |mapping| mapping.status.is_none() && mapping.error_str.contains("invalid type:"), + PublicError::BadRequest, + ), + Rule::new( + |mapping| mapping.status.is_none() && mapping.error_str.contains("Unexpected server error"), + PublicError::InternalServer, + ), ]; -pub(super) fn map(mapping: &Mapping<'_>) -> Option { - apply(RULES, mapping).map(|failure| PublicFailure { - llm_provider: Some("cohere".to_string()), - ..failure - }) -} - #[cfg(test)] mod tests { - use super::super::testing::{context, failure, http, status, upstream}; - use super::super::{ExceptionFamily, OriginalException, PublicKind}; + use super::super::rules::first_match; + use super::super::testing::mapping; use super::*; - fn mapped(provider: &str, original: &OriginalException) -> Option { - let context = context(provider, ExceptionFamily::Cohere); - map(&Mapping::new(&context, original)) + fn classified(text: &str) -> Option { + classified_with(Some(400), text) } - fn cohere(class: StatusClass, status_code: u16, body: &str, message: &str) -> PublicFailure { - failure( - status(class, upstream(status_code, body)), - message, - "cohere", - ) + fn classified_with(status: Option, text: &str) -> Option { + first_match(RULES, &mapping(status, text)).map(|rule| rule.error) } #[rstest::rstest] - #[case::invalid_token( - 500, - "invalid api token", - cohere( - StatusClass::Authentication, - 500, - "invalid api token", - "CohereException - invalid api token" - ) - )] - #[case::no_api_key( - 500, - "No API key provided.", - cohere( - StatusClass::Authentication, - 500, - "No API key provided.", - "CohereException - No API key provided." - ) - )] - #[case::invalid_parameter( - 500, - "invalid type: parameter x", - cohere( - StatusClass::BadRequest, - 500, - "invalid type: parameter x", - "CohereException - invalid type: parameter x" - ) - )] - #[case::too_many_tokens( - 500, - "too many tokens", - cohere( - StatusClass::ContextWindowExceeded, - 500, - "too many tokens", - "CohereException - too many tokens" - ) - )] - #[case::internal_server_text( - 400, - "Internal Server Error", - cohere( - StatusClass::InternalServer, - 400, - "Internal Server Error", - "CohereException - Internal Server Error" - ) - )] - #[case::bad_request( - 400, - "rejected", - cohere(StatusClass::BadRequest, 400, "rejected", "CohereException - rejected") - )] - #[case::invalid_token_status( - 498, - "rejected", - cohere(StatusClass::BadRequest, 498, "rejected", "CohereException - rejected") - )] - #[case::request_timeout(408, "rejected", failure(PublicKind::Timeout { status: None }, "CohereException - rejected", "cohere"))] - #[case::internal_server( - 500, - "rejected", - cohere( - StatusClass::InternalServer, - 500, - "rejected", - "CohereException - rejected" - ) - )] - fn each_rule_maps_and_reports_cohere( - #[case] status_code: u16, - #[case] body: &str, - #[case] expected: PublicFailure, - ) { - assert_eq!(mapped("azure_ai", &http(status_code, body)), Some(expected)); - } - - #[rstest::rstest] - #[case::unmapped_status(409)] - #[case::unauthorized(401)] - fn statuses_without_a_rule_fall_through(#[case] status_code: u16) { - assert_eq!(mapped("cohere", &http(status_code, "rejected")), None); - } - - #[test] - fn the_internal_server_rule_uses_the_redacted_text() { - let body = "internal server error Bearer abcdefghijklmnop"; - assert_eq!( - mapped("cohere", &http(400, body)), - Some(cohere( - StatusClass::InternalServer, - 400, - body, - "CohereException - internal server error REDACTED" - )) - ); + #[case::invalid_token("invalid api token", PublicError::Authentication)] + #[case::no_api_key("No API key provided.", PublicError::Authentication)] + #[case::invalid_parameter("invalid type: parameter x", PublicError::BadRequest)] + #[case::too_many_tokens("too many tokens", PublicError::ContextWindowExceeded)] + #[case::internal_server_text("Internal Server Error", PublicError::InternalServer)] + #[case::internal_server_any_case("INTERNAL server ERROR", PublicError::InternalServer)] + fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) { + assert_eq!(classified(text), Some(expected)); } #[rstest::rstest] #[case::token_before_parameter( "invalid api token invalid type: parameter", - StatusClass::Authentication + PublicError::Authentication )] #[case::parameter_before_tokens( "invalid type: parameter too many tokens", - StatusClass::BadRequest + PublicError::BadRequest )] #[case::tokens_before_internal( "too many tokens Internal Server Error", - StatusClass::ContextWindowExceeded + PublicError::ContextWindowExceeded )] - #[case::internal_before_status("Internal Server Error", StatusClass::InternalServer)] - fn the_earlier_rule_wins_when_two_apply(#[case] body: &str, #[case] class: StatusClass) { - assert_eq!( - mapped("cohere", &http(400, body)), - Some(cohere( - class, - 400, - body, - &format!("CohereException - {body}") - )) - ); + fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) { + assert_eq!(classified(text), Some(expected)); + } + + #[rstest::rstest] + #[case::invalid_type(None, "invalid type: x", Some(PublicError::BadRequest))] + #[case::unexpected_server_error( + None, + "Unexpected server error", + Some(PublicError::InternalServer) + )] + #[case::invalid_type_before_unexpected( + None, + "invalid type: x Unexpected server error", + Some(PublicError::BadRequest) + )] + #[case::internal_before_invalid_type( + None, + "internal server error invalid type: x", + Some(PublicError::InternalServer) + )] + #[case::invalid_type_with_a_status(Some(500), "invalid type: x", None)] + #[case::unexpected_with_a_status(Some(400), "Unexpected server error", None)] + fn the_trailing_rules_only_claim_failures_without_a_status( + #[case] status: Option, + #[case] text: &str, + #[case] expected: Option, + ) { + assert_eq!(classified_with(status, text), expected); + } + + #[test] + fn text_without_a_marker_is_left_to_the_status_table() { + assert_eq!(classified("rejected"), None); } } diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs index acc7820c770..162d325e4f4 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs @@ -1,18 +1,47 @@ -//! A port of Python's `exception_type` for the routes that run in Rust. +//! A port of Python's `exception_type` for the routes that run in Rust. Rust decides the +//! public class, the message and the debug text; Python only builds the class. +//! +//! DIVERGENCES: where the Python mapper is inconsistent, the port follows one rule instead. +//! - The message is always `{Provider}Exception - {redacted text}`. Python's per-branch +//! labels (`RateLimitError: `, `litellm.RateLimitError: `, `Vertex_aiException BadRequestError`) +//! are dropped because every public class already prefixes `litellm.{Class}: `. +//! - The upstream response is always the real one. Python swaps in made-up `httpx.Response` +//! stubs on some Vertex branches, losing the body and `retry-after`. +//! - The debug text is always attached; Python passes it on some branches only. +//! - No family rule turns a status into a class; the shared status table owns that. So a +//! Vertex 502 is a `BadGatewayError` and an OpenAI-family 403 is a `PermissionDeniedError`. +//! Three rules read the status only to gate a text match, as Python does: the standalone +//! `429`, Vertex's wrapped 429 behind a 5xx, and Cohere's rules for failures with no status. +//! - A timeout text marker on an HTTP failure keeps the upstream response. Python's `Timeout` +//! carries none. +//! - Every family matches and reports the redacted text. Python's OpenAI mapper builds the +//! message from the unredacted text. +//! - A refused connection is an `APIConnectionError`, not the 500 Python's HTTP handler +//! synthesizes. +//! - Dropped Python rules: Vertex's bare `403` substring (it matches `4031 tokens`), Vertex's +//! `IndexError` quota marker (a Python client crash), the OpenAI SDK's missing-`api_key` +//! text and its `OPENAI` renaming, Cohere's `llm_provider="cohere"` override, and Cohere's +//! `CohereConnectionError` check (a Python SDK class name). //! //! KNOWN_GAPS: differences from the Python mapper that no Rust route can reach today. Each //! one stops being acceptable at its trigger. -//! - The Vertex partner-model API base for "claude" models is not built into -//! `extra_information`. Trigger: a Vertex route whose models include Anthropic partner -//! models; then `api_base` gets that branch and a table row. +//! - The Vertex partner-model API base for "claude" models is not built into the debug text. +//! Trigger: a Vertex route whose models include Anthropic partner models. +//! - The debug text's `API Base` line is only the non-streaming Vertex URL. Python prefers an +//! explicit or provider-resolved `api_base`, uses `:streamGenerateContent` when streaming, +//! and has Gemini and OpenAI defaults. Trigger: the first route wired to this mapper, since +//! every route knows its `api_base`. +//! - The debug text has no `Messages:` line, which Python adds when +//! `redact_messages_in_exceptions` is off. Trigger: a wired route that carries messages. //! - Python reports the provider `get_llm_provider` resolves for a stripped model name when //! that name happens to be in the model cost map. Trigger: a route whose model names //! overlap the cost map; that needs the provider resolution port, not a classifier change. -//! - The generic `APIConnectionError` fallback appends `traceback.format_exc()` to the -//! message. Rust has no Python traceback and does not invent one; a sweep row that reaches -//! it compares the message before the traceback. +//! - `litellm_proxy` errors are not unwrapped into the proxied exception. Trigger: a Rust +//! route that calls a LiteLLM proxy. +//! - Only the OpenAI-compatible, Vertex AI and Cohere mappers are ported; every other +//! provider goes straight to the status table. Trigger: a Rust route for such a provider. -use super::secret_redaction::{redact_string, secret_redaction_enabled}; +use super::secret_redaction::SecretRedactor; mod cohere; mod openai; @@ -22,10 +51,10 @@ mod rules; mod status; mod vertex_ai; -pub use original::{ExceptionFamily, LocalClass, OriginalException}; -pub use public::{HttpStub, PublicFailure, PublicKind, ResponseArg, StatusClass, UpstreamResponse}; +pub use original::{ExceptionFamily, OriginalException}; +pub use public::{MappedFailure, PublicError, UpstreamResponse}; -const DOCS_URL: &str = "https://docs.litellm.ai/docs"; +use rules::{Rule, contains_any, first_match}; const TIMEOUT_MARKERS: &[&str] = &[ "Request Timeout Error", @@ -37,11 +66,8 @@ const TIMEOUT_MARKERS: &[&str] = &[ #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct ExceptionContext { pub model: String, - pub custom_llm_provider: Option, - pub family: ExceptionFamily, + pub custom_llm_provider: String, pub asynchronous: bool, - pub suppress_debug_info: bool, - pub redact_messages_in_exceptions: bool, pub vertex_project: Option, pub vertex_location: Option, pub model_group: Option, @@ -50,73 +76,93 @@ pub struct ExceptionContext { pub user_api_key_team_alias: Option, } -/// The attributes `exception_type` reads off the Python exception: a provider error -/// (`BaseLLMException`) carries a status, a response and a request, a plain exception -/// carries only its text. -struct Raised { +/// What the rules read: the status of a provider response, if any, and the redacted text. +struct Mapping { status: Option, - status_is_synthesized: bool, - message: String, - response: Option, + error_str: String, } -impl Raised { - fn provider( - status: u16, - message: String, - body: String, - headers: Vec<(String, String)>, - ) -> Self { - Self { - status: Some(status), - status_is_synthesized: false, - message, - response: Some(UpstreamResponse { - status, - body, - headers, +pub fn exception_type( + context: &ExceptionContext, + redactor: Option<&SecretRedactor>, + original: &OriginalException, +) -> MappedFailure { + let (status, text, upstream) = match original { + OriginalException::Http { + status, + body, + headers, + } => ( + Some(*status), + body.clone(), + Some(UpstreamResponse { + status: *status, + body: body.clone(), + headers: headers.clone(), }), + ), + OriginalException::Connection { message } | OriginalException::Plain { message } => { + (None, message.clone(), None) } - } - - fn plain(message: String) -> Self { - Self { - status: None, - status_is_synthesized: false, - message, - response: None, - } - } - - fn new(original: &OriginalException, asynchronous: bool) -> Self { - match original { - OriginalException::Http { - status, - body, - headers, - } => Self::provider(*status, body.clone(), body.clone(), headers.clone()), - OriginalException::Connection { message } => Self { - status_is_synthesized: true, - ..Self::provider(500, message.clone(), String::new(), Vec::new()) - }, - OriginalException::Timeout { - timeout_seconds, - elapsed_seconds, - } => Self::provider( - 408, - timeout_message(asynchronous, *timeout_seconds, *elapsed_seconds), - String::new(), - Vec::new(), - ), - OriginalException::Response { message } - | OriginalException::Local { message, .. } - | OriginalException::Public { message, .. } => Self::plain(message.clone()), - } + OriginalException::Timeout { + timeout_seconds, + elapsed_seconds, + } => ( + None, + timeout_message(context.asynchronous, *timeout_seconds, *elapsed_seconds), + None, + ), + }; + let mapping = Mapping { + status, + error_str: match redactor { + Some(redactor) => redactor.redact(&text), + None => text, + }, + }; + let family = ExceptionFamily::for_provider(&context.custom_llm_provider); + let (error, hint) = classify(family, original, &mapping); + MappedFailure { + error, + message: format!( + "{} - {}{hint}", + exception_provider(&context.custom_llm_provider), + mapping.error_str + ), + upstream, + debug_info: extra_information(context, api_base(context).as_deref()), } } -/// The text `litellm.Timeout` carries when the Python HTTP handler times out: the sync -/// and async handlers word it differently. +fn classify( + family: ExceptionFamily, + original: &OriginalException, + mapping: &Mapping, +) -> (PublicError, &'static str) { + const TIMEOUT: PublicError = PublicError::Timeout { status: 408 }; + if matches!(original, OriginalException::Timeout { .. }) + || contains_any(&mapping.error_str, TIMEOUT_MARKERS) + { + return (TIMEOUT, ""); + } + if let Some(rule) = first_match(family_rules(family), mapping) { + return (rule.error, rule.hint); + } + let by_status = mapping.status.and_then(status::classify); + (by_status.unwrap_or(PublicError::ApiConnection), "") +} + +fn family_rules(family: ExceptionFamily) -> &'static [Rule] { + match family { + ExceptionFamily::OpenAiCompatible => openai::RULES, + ExceptionFamily::VertexAi => vertex_ai::RULES, + ExceptionFamily::Cohere => cohere::RULES, + ExceptionFamily::Other => &[], + } +} + +/// The text the Python HTTP handler's timeout carries: the sync and async handlers word it +/// differently. fn timeout_message( asynchronous: bool, timeout_seconds: Option, @@ -126,11 +172,9 @@ fn timeout_message( if asynchronous { let elapsed = python_float(elapsed_seconds.map(|seconds| (seconds * 1000.0).round() / 1000.0)); - format!( - "litellm.Timeout: Connection timed out. Timeout passed={timeout}, time taken={elapsed} seconds" - ) + format!("Connection timed out. Timeout passed={timeout}, time taken={elapsed} seconds") } else { - format!("litellm.Timeout: Connection timed out after {timeout} seconds.") + format!("Connection timed out after {timeout} seconds.") } } @@ -142,113 +186,10 @@ fn python_float(value: Option) -> String { } } -/// Everything the rules read: the original as Python sees it and the text `exception_type` -/// derives from the context before any provider mapper runs. -struct Mapping<'a> { - context: &'a ExceptionContext, - original: Raised, - provider: &'a str, - error_str: String, - exception_provider: String, - extra_information: String, -} - -impl<'a> Mapping<'a> { - fn new(context: &'a ExceptionContext, original: &OriginalException) -> Self { - let original = Raised::new(original, context.asynchronous); - let error_str = if secret_redaction_enabled() { - redact_string(&original.message) - } else { - original.message.clone() - }; - Self { - context, - original, - provider: context.custom_llm_provider.as_deref().unwrap_or_default(), - error_str, - exception_provider: match &context.custom_llm_provider { - None => "None".to_string(), - Some(provider) => exception_provider(provider), - }, - extra_information: extra_information(context, api_base(context).as_deref()), - } - } - - fn failure(&self, kind: PublicKind, message: String, debug: bool) -> PublicFailure { - PublicFailure { - kind, - message, - model: self.context.model.clone(), - llm_provider: self.context.custom_llm_provider.clone(), - litellm_debug_info: debug.then(|| self.extra_information.clone()), - litellm_response_headers: None, - print_banner: false, - } - } -} - -pub fn exception_type(context: &ExceptionContext, original: &OriginalException) -> PublicFailure { - if let OriginalException::Public { class, message } = original { - return PublicFailure { - kind: PublicKind::Status { - status_class: *class, - response: None, - }, - message: message.clone(), - model: context.model.clone(), - llm_provider: context.custom_llm_provider.clone(), - litellm_debug_info: None, - litellm_response_headers: None, - print_banner: false, - }; - } - let mapping = Mapping::new(context, original); - let litellm_response_headers = mapping - .original - .response - .as_ref() - .map(|response| response.headers.clone()) - .filter(|headers| !headers.is_empty()); - PublicFailure { - litellm_response_headers, - print_banner: !context.suppress_debug_info, - ..map(&mapping) - } -} - -fn map(mapping: &Mapping<'_>) -> PublicFailure { - if rules::contains_any(&mapping.error_str, TIMEOUT_MARKERS) { - return mapping.failure( - PublicKind::Timeout { status: None }, - format!( - "APITimeoutError - Request timed out. Error_str: {}", - mapping.error_str - ), - true, - ); - } - let provider_failure = match mapping.context.family { - ExceptionFamily::OpenAiCompatible => openai::map(mapping), - ExceptionFamily::VertexAi => vertex_ai::map(mapping), - ExceptionFamily::Cohere => cohere::map(mapping), - ExceptionFamily::Other => None, - }; - provider_failure - .or_else(|| status::map(mapping)) - .unwrap_or_else(|| unmapped(mapping)) -} - -/// The `APIConnectionError` Python raises when no mapper claimed the failure: with the -/// provider prefix for a provider error, with the bare text for a plain exception. -fn unmapped(mapping: &Mapping<'_>) -> PublicFailure { - let message = match mapping.original.status { - Some(_) => format!("{} - {}", mapping.exception_provider, mapping.error_str), - None => mapping.original.message.clone(), - }; - mapping.failure(PublicKind::ApiConnection, message, false) -} - fn exception_provider(provider: &str) -> String { + if provider == "openai" { + return "OpenAIException".to_string(); + } let mut characters = provider.chars(); match characters.next() { Some(first) => format!("{}{}Exception", first.to_uppercase(), characters.as_str()), @@ -256,18 +197,6 @@ fn exception_provider(provider: &str) -> String { } } -fn python_capitalize(value: &str) -> String { - let mut characters = value.chars(); - match characters.next() { - Some(first) => format!( - "{}{}", - first.to_uppercase(), - characters.as_str().to_lowercase() - ), - None => String::new(), - } -} - fn api_base(context: &ExceptionContext) -> Option { match (&context.vertex_location, &context.vertex_project) { (Some(location), Some(project)) => Some(format!( @@ -311,132 +240,99 @@ fn extra_information(context: &ExceptionContext, api_base: Option<&str>) -> Stri #[cfg(test)] mod testing { - use super::*; + use super::Mapping; - pub(super) const DEBUG: &str = "\nModel: ocr-model"; - - pub(super) fn context(provider: &str, family: ExceptionFamily) -> ExceptionContext { - ExceptionContext { - model: "ocr-model".into(), - custom_llm_provider: Some(provider.into()), - family, - suppress_debug_info: true, - ..ExceptionContext::default() - } - } - - pub(super) fn http(status: u16, body: &str) -> OriginalException { - OriginalException::Http { + pub(super) fn mapping(status: Option, text: &str) -> Mapping { + Mapping { status, - body: body.into(), - headers: vec![("retry-after".into(), "7".into())], - } - } - - pub(super) fn upstream(status: u16, body: &str) -> Option { - Some(ResponseArg::Upstream(UpstreamResponse { - status, - body: body.into(), - headers: vec![("retry-after".into(), "7".into())], - })) - } - - pub(super) fn status(class: StatusClass, response: Option) -> PublicKind { - PublicKind::Status { - status_class: class, - response, - } - } - - /// The failure a rule builds before `exception_type` adds the response headers and the - /// banner flag. - pub(super) fn failure(kind: PublicKind, message: &str, provider: &str) -> PublicFailure { - PublicFailure { - kind, - message: message.into(), - model: "ocr-model".into(), - llm_provider: Some(provider.into()), - litellm_debug_info: None, - litellm_response_headers: None, - print_banner: false, - } - } - - pub(super) fn with_debug(failure: PublicFailure) -> PublicFailure { - PublicFailure { - litellm_debug_info: Some(DEBUG.into()), - ..failure + error_str: text.into(), } } } #[cfg(test)] mod tests { - use super::testing::{DEBUG, context, failure, http, status, upstream, with_debug}; use super::*; - fn openai() -> ExceptionContext { - context("mistral", ExceptionFamily::OpenAiCompatible) + const DEBUG: &str = "\nModel: ocr-model"; + + fn context(provider: &str) -> ExceptionContext { + ExceptionContext { + model: "ocr-model".into(), + custom_llm_provider: provider.into(), + ..ExceptionContext::default() + } + } + + fn redactor() -> SecretRedactor { + SecretRedactor::new(16) + } + + fn headers() -> Vec<(String, String)> { + vec![("retry-after".into(), "7".into())] + } + + fn http(status: u16, body: &str) -> OriginalException { + OriginalException::Http { + status, + body: body.into(), + headers: headers(), + } + } + + fn upstream(status: u16, body: &str) -> Option { + Some(UpstreamResponse { + status, + body: body.into(), + headers: headers(), + }) + } + + fn mapped(provider: &str, original: &OriginalException) -> MappedFailure { + exception_type(&context(provider), Some(&redactor()), original) + } + + #[rstest::rstest] + #[case::openai_family("mistral", "rate limit reached", PublicError::RateLimit)] + #[case::vertex_family("vertex_ai", "Resource exhausted", PublicError::RateLimit)] + #[case::cohere_family("cohere", "too many tokens", PublicError::ContextWindowExceeded)] + fn a_family_text_rule_beats_the_status_and_keeps_the_real_response( + #[case] provider: &str, + #[case] body: &str, + #[case] expected: PublicError, + ) { + let failure = mapped(provider, &http(401, body)); + assert_eq!(failure.error, expected); + assert_eq!(failure.upstream, upstream(401, body)); } #[test] - fn a_public_original_passes_through_without_banner_debug_or_prefix() { - let original = OriginalException::Public { - class: StatusClass::UnsupportedParams, - message: "Invalid `req_format`".into(), - }; - let context = ExceptionContext { - suppress_debug_info: false, - ..openai() - }; + fn the_other_family_has_no_text_rules() { assert_eq!( - exception_type(&context, &original), - failure( - status(StatusClass::UnsupportedParams, None), - "Invalid `req_format`", - "mistral" - ) + mapped("reducto", &http(401, "rate limit reached")).error, + PublicError::Authentication ); } #[rstest::rstest] - #[case::vertex_family_status_rule(ExceptionFamily::VertexAi, "vertex_ai", PublicFailure { - litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), - ..with_debug(failure(status(StatusClass::BadRequest, upstream(409, "rejected")), "Vertex_aiException - rejected", "vertex_ai")) - })] - #[case::cohere_family(ExceptionFamily::Cohere, "cohere", PublicFailure { - litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), - ..with_debug(failure(status(StatusClass::BadRequest, upstream(409, "rejected")), "CohereException - rejected", "cohere")) - })] - #[case::other_family(ExceptionFamily::Other, "reducto", PublicFailure { - litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), - ..with_debug(failure(status(StatusClass::BadRequest, upstream(409, "rejected")), "ReductoException - rejected", "reducto")) - })] - fn families_without_a_409_rule_reach_the_status_table( - #[case] family: ExceptionFamily, + #[case::openai_403_is_permission_denied("mistral", 403, PublicError::PermissionDenied)] + #[case::openai_409_is_bad_request("mistral", 409, PublicError::BadRequest)] + #[case::vertex_502_is_bad_gateway("vertex_ai", 502, PublicError::BadGateway)] + #[case::vertex_504_is_a_timeout("vertex_ai", 504, PublicError::Timeout { status: 504 })] + #[case::cohere_498_is_bad_request("cohere", 498, PublicError::BadRequest)] + #[case::other_503("reducto", 503, PublicError::ServiceUnavailable)] + fn without_a_text_rule_every_family_uses_the_status_table( #[case] provider: &str, - #[case] expected: PublicFailure, + #[case] status: u16, + #[case] expected: PublicError, ) { assert_eq!( - exception_type(&context(provider, family), &http(409, "rejected")), - expected - ); - } - - #[test] - fn the_openai_family_claims_a_409_before_the_status_table() { - assert_eq!( - exception_type(&openai(), &http(409, "rejected")), - PublicFailure { - litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), - ..with_debug(failure( - PublicKind::Api { - status: 409, - request_url: DOCS_URL - }, - "APIError: MistralException - rejected", - "mistral" - )) + mapped(provider, &http(status, "rejected")), + MappedFailure { + error: expected, + message: format!("{} - rejected", exception_provider(provider)), + upstream: upstream(status, "rejected"), + debug_info: DEBUG.into(), } ); } @@ -447,148 +343,126 @@ mod tests { #[case::timed_out_generating("Timed out generating response")] #[case::read_operation("The read operation timed out")] fn timeout_markers_win_over_every_family(#[case] marker: &str) { - let body = format!("rate limit {marker}"); - for family in [ - ExceptionFamily::OpenAiCompatible, - ExceptionFamily::VertexAi, - ExceptionFamily::Cohere, - ExceptionFamily::Other, - ] { + let body = format!("rate limit invalid api token {marker}"); + for provider in ["mistral", "vertex_ai", "cohere", "reducto"] { assert_eq!( - exception_type(&context("mistral", family), &http(429, &body)), - PublicFailure { - litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), - ..with_debug(failure( - PublicKind::Timeout { status: None }, - &format!("APITimeoutError - Request timed out. Error_str: {body}"), - "mistral" - )) - } + mapped(provider, &http(429, &body)).error, + PublicError::Timeout { status: 408 }, + "{provider}" ); } } + #[test] + fn a_handler_timeout_is_a_408_without_a_response() { + let original = OriginalException::Timeout { + timeout_seconds: Some(0.5), + elapsed_seconds: Some(0.5031), + }; + assert_eq!( + mapped("mistral", &original), + MappedFailure { + error: PublicError::Timeout { status: 408 }, + message: "MistralException - Connection timed out after 0.5 seconds.".into(), + upstream: None, + debug_info: DEBUG.into(), + } + ); + } + #[rstest::rstest] - #[case::provider_error_keeps_the_prefix(http(409, "rejected"), "ReductoException - rejected")] - #[case::synthesized_status_skips_the_status_table( - OriginalException::Connection { message: "refused".into() }, - "ReductoException - refused" - )] - #[case::plain_exception_keeps_its_text( - OriginalException::Local { class: LocalClass::FileNotFound, message: "File not found: /a".into() }, - "File not found: /a" - )] - fn unmapped_failures_are_connection_errors( + #[case::refused_connection(OriginalException::Connection { message: "refused".into() })] + #[case::unparseable_response(OriginalException::Plain { message: "refused".into() })] + #[case::informational_status(OriginalException::Http { status: 399, body: "refused".into(), headers: Vec::new() })] + fn a_failure_no_rule_or_status_claims_is_a_connection_error( #[case] original: OriginalException, - #[case] message: &str, ) { - let context = context("reducto", ExceptionFamily::Other); - let expected = match original { - OriginalException::Http { .. } => with_debug(failure( - status(StatusClass::BadRequest, upstream(409, "rejected")), - message, - "reducto", - )), - OriginalException::Connection { .. } | OriginalException::Local { .. } => { - failure(PublicKind::ApiConnection, message, "reducto") - } - _ => unreachable!(), - }; - let actual = exception_type(&context, &original); - assert_eq!( - PublicFailure { - litellm_response_headers: None, - ..actual - }, - expected - ); + let failure = mapped("reducto", &original); + assert_eq!(failure.error, PublicError::ApiConnection); + assert_eq!(failure.message, "ReductoException - refused"); } #[test] - fn a_missing_provider_renders_like_python_none() { - let context = ExceptionContext { - custom_llm_provider: None, - family: ExceptionFamily::Other, - ..openai() + fn a_timeout_marker_on_a_response_keeps_the_response() { + let failure = mapped("reducto", &http(429, "Request timed out")); + assert_eq!(failure.error, PublicError::Timeout { status: 408 }); + assert_eq!(failure.upstream, upstream(429, "Request timed out")); + } + + #[test] + fn family_text_rules_also_classify_failures_without_a_response() { + let original = OriginalException::Plain { + message: "Request too large".into(), }; - assert_eq!( - exception_type(&context, &http(401, "rejected")), - PublicFailure { - llm_provider: None, - litellm_response_headers: Some(vec![("retry-after".into(), "7".into())]), - ..with_debug(failure( - status(StatusClass::Authentication, upstream(401, "rejected")), - "None - rejected", - "unused" - )) - } - ); + assert_eq!(mapped("mistral", &original).error, PublicError::RateLimit); } #[rstest::rstest] - #[case::suppressed(true, false)] - #[case::printed(false, true)] - fn the_banner_prints_unless_debug_info_is_suppressed( - #[case] suppress_debug_info: bool, - #[case] print_banner: bool, - ) { - let context = ExceptionContext { - suppress_debug_info, - ..openai() - }; + #[case::openai_family("mistral", "MistralException - rejected REDACTED")] + #[case::vertex_family("vertex_ai", "Vertex_aiException - rejected REDACTED")] + #[case::other_family("reducto", "ReductoException - rejected REDACTED")] + fn every_family_reports_the_redacted_text(#[case] provider: &str, #[case] message: &str) { + let failure = mapped(provider, &http(400, "rejected Bearer abcdefghijklmnop")); + assert_eq!(failure.message, message); + } + + #[test] + fn redaction_runs_before_the_rules_see_the_text() { + let body = "db_password=rate_limit"; assert_eq!( - exception_type(&context, &http(400, "rejected")).print_banner, - print_banner + mapped("mistral", &http(400, body)).error, + PublicError::BadRequest + ); + assert_eq!( + exception_type(&context("mistral"), None, &http(400, body)).error, + PublicError::RateLimit ); } #[test] - fn empty_upstream_headers_are_not_reported() { - let original = OriginalException::Http { - status: 400, - body: "rejected".into(), - headers: Vec::new(), - }; - assert_eq!( - exception_type(&openai(), &original).litellm_response_headers, - None - ); - } - - #[test] - fn messages_are_redacted_before_markers_and_prefixes() { + fn without_a_redactor_the_text_is_kept() { let body = "rejected Bearer abcdefghijklmnop"; assert_eq!( - exception_type( - &context("reducto", ExceptionFamily::Other), - &http(400, body) - ) - .message, - "ReductoException - rejected REDACTED" + exception_type(&context("reducto"), None, &http(400, body)).message, + format!("ReductoException - {body}") ); } - const SYNC_TIMEOUT: &str = "litellm.Timeout: Connection timed out after 0.5 seconds."; + #[test] + fn a_rule_hint_follows_the_message() { + let failure = mapped("mistral", &http(400, "invalid_encrypted_content")); + assert_eq!(failure.error, PublicError::BadRequest); + assert!( + failure + .message + .starts_with("MistralException - invalid_encrypted_content\n\n This error occurs") + ); + } #[rstest::rstest] - #[case::sync(false, Some(0.5), Some(0.5031), SYNC_TIMEOUT)] + #[case::sync( + false, + Some(0.5), + Some(0.5031), + "Connection timed out after 0.5 seconds." + )] #[case::async_rounds_the_elapsed_time( true, Some(0.5), Some(0.5031), - "litellm.Timeout: Connection timed out. Timeout passed=0.5, time taken=0.503 seconds" + "Connection timed out. Timeout passed=0.5, time taken=0.503 seconds" )] #[case::whole_seconds_keep_a_decimal( true, Some(600.0), Some(2.0), - "litellm.Timeout: Connection timed out. Timeout passed=600.0, time taken=2.0 seconds" + "Connection timed out. Timeout passed=600.0, time taken=2.0 seconds" )] #[case::unknown_values_render_as_none( true, None, None, - "litellm.Timeout: Connection timed out. Timeout passed=None, time taken=None seconds" + "Connection timed out. Timeout passed=None, time taken=None seconds" )] fn timeout_text_follows_the_delivery_mode( #[case] asynchronous: bool, @@ -602,58 +476,6 @@ mod tests { ); } - #[rstest::rstest] - #[case::sync(false, SYNC_TIMEOUT)] - #[case::async_( - true, - "litellm.Timeout: Connection timed out. Timeout passed=0.5, time taken=0.503 seconds" - )] - fn a_timeout_is_a_408_carrying_the_handler_text( - #[case] asynchronous: bool, - #[case] text: &str, - ) { - let context = ExceptionContext { - asynchronous, - ..openai() - }; - let original = OriginalException::Timeout { - timeout_seconds: Some(0.5), - elapsed_seconds: Some(0.5031), - }; - assert_eq!( - exception_type(&context, &original), - with_debug(failure( - PublicKind::Timeout { status: None }, - &format!("Timeout Error: MistralException - {text}"), - "mistral" - )) - ); - } - - #[test] - fn a_refused_connection_is_a_500_with_an_empty_response() { - assert_eq!( - exception_type( - &openai(), - &OriginalException::Connection { - message: "refused".into() - } - ), - with_debug(failure( - status( - StatusClass::InternalServer, - Some(ResponseArg::Upstream(UpstreamResponse { - status: 500, - body: String::new(), - headers: Vec::new(), - })) - ), - "InternalServerError: MistralException - refused", - "mistral" - )) - ); - } - #[test] fn debug_information_follows_the_python_layout() { let context = ExceptionContext { @@ -662,10 +484,10 @@ mod tests { model_group: Some("ocr".into()), deployment: Some("deployment".into()), user_api_key_alias: Some("key".into()), - ..openai() + ..context("vertex_ai") }; assert_eq!( - extra_information(&context, api_base(&context).as_deref()), + exception_type(&context, None, &http(400, "rejected")).debug_info, concat!( "\n\nKey Name: `key`\nTeam: `None`", "\nModel: ocr-model", @@ -680,21 +502,20 @@ mod tests { #[rstest::rstest] #[case::bare(ExceptionContext::default(), "\nModel: ")] - #[case::redacted_messages(ExceptionContext { redact_messages_in_exceptions: true, model: "m".into(), ..ExceptionContext::default() }, "\nModel: m")] #[case::team_alias( - ExceptionContext { model: "m".into(), user_api_key_alias: Some("key".into()), user_api_key_team_alias: Some("team".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, + ExceptionContext { model: "m".into(), user_api_key_alias: Some("key".into()), user_api_key_team_alias: Some("team".into()), ..ExceptionContext::default() }, "\n\nKey Name: `key`\nTeam: `team`\nModel: m" )] #[case::team_alias_without_key_is_ignored( - ExceptionContext { model: "m".into(), user_api_key_team_alias: Some("team".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, + ExceptionContext { model: "m".into(), user_api_key_team_alias: Some("team".into()), ..ExceptionContext::default() }, "\nModel: m" )] #[case::project_without_location_has_no_api_base( - ExceptionContext { model: "m".into(), vertex_project: Some("p".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, + ExceptionContext { model: "m".into(), vertex_project: Some("p".into()), ..ExceptionContext::default() }, "\nModel: m\nvertex_project: `p`\n" )] #[case::location_without_project_has_no_api_base( - ExceptionContext { model: "m".into(), vertex_location: Some("l".into()), redact_messages_in_exceptions: true, ..ExceptionContext::default() }, + ExceptionContext { model: "m".into(), vertex_location: Some("l".into()), ..ExceptionContext::default() }, "\nModel: m\nvertex_location: `l`\n" )] fn each_optional_context_field_adds_its_own_line( @@ -708,6 +529,7 @@ mod tests { } #[rstest::rstest] + #[case::openai_keeps_its_brand("openai", "OpenAIException")] #[case::lowercase("mistral", "MistralException")] #[case::keeps_the_rest("azure_ai", "Azure_aiException")] #[case::empty("", "")] @@ -717,17 +539,4 @@ mod tests { ) { assert_eq!(exception_provider(provider), expected); } - - #[rstest::rstest] - #[case::lowers_the_rest("vERTEX_AI", "Vertex_ai")] - #[case::empty("", "")] - fn python_capitalize_lowers_the_rest(#[case] value: &str, #[case] expected: &str) { - assert_eq!(python_capitalize(value), expected); - } - - #[test] - fn debug_constant_matches_the_default_test_context() { - let context = openai(); - assert_eq!(extra_information(&context, None), DEBUG); - } } diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs index ffbb172582b..b4f4496dbf0 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs @@ -1,85 +1,31 @@ -use super::public::{PublicFailure, StatusClass}; -use super::rules::{ - ApiStatus, Kind, ResponseChoice, Rule, apply, contains_any, is_context_window_exceeded, - is_rate_limit, -}; -use super::{DOCS_URL, Mapping}; - -const OPENAI_URL: &str = "https://api.openai.com/v1"; +use super::public::PublicError; +use super::rules::{Rule, contains_any, is_context_window_exceeded, is_rate_limit}; const ENCRYPTED_CONTENT_HELP: &str = "\n\n This error occurs when load balancing Responses API across deployments with different API keys.\n Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n router_settings:\n enable_pre_call_checks: true\n optional_pre_call_checks:\n - encrypted_content_affinity\n\n Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing"; -const fn with_response(class: StatusClass) -> Kind { - Kind::Status { - class, - response: ResponseChoice::Provider, - } -} - -fn exception_provider(mapping: &Mapping<'_>) -> String { - if mapping.provider == "openai" { - "OpenAIException".to_string() - } else { - super::exception_provider(mapping.provider) - } -} - -/// The raw message with OpenAI's own names swapped for the provider's. -fn message(mapping: &Mapping<'_>) -> String { - let provider = mapping.provider; - mapping - .original - .message - .replace("OPENAI", &provider.to_uppercase()) - .replace("openai.OpenAIError", &format!("{provider}.{provider}Error")) -} - -fn prefixed(mapping: &Mapping<'_>, label: &str) -> String { - format!( - "{label}{} - {}", - exception_provider(mapping), - message(mapping) - ) -} - -fn status_is(mapping: &Mapping<'_>, statuses: &[u16]) -> bool { - mapping - .original - .status - .is_some_and(|status| statuses.contains(&status)) -} - -/// `_map_openai_exception`, in its branch order. -const RULES: &[Rule] = &[ - Rule { - when: |mapping| is_rate_limit(&mapping.error_str, mapping.original.status), - kind: with_response(StatusClass::RateLimit), - message: |mapping| prefixed(mapping, "RateLimitError: "), - debug: false, - }, - Rule { - when: |mapping| is_context_window_exceeded(&mapping.error_str), - kind: with_response(StatusClass::ContextWindowExceeded), - message: |mapping| prefixed(mapping, "ContextWindowExceededError: "), - debug: true, - }, - Rule { - when: |mapping| { +/// The text branches of `_map_openai_exception`, in its order. +pub(super) const RULES: &[Rule] = &[ + Rule::new( + |mapping| is_rate_limit(&mapping.error_str, mapping.status), + PublicError::RateLimit, + ), + Rule::new( + |mapping| is_context_window_exceeded(&mapping.error_str), + PublicError::ContextWindowExceeded, + ), + Rule::new( + |mapping| { mapping.error_str.contains("invalid_request_error") && mapping.error_str.contains("model_not_found") }, - kind: with_response(StatusClass::NotFound), - message: |mapping| prefixed(mapping, ""), - debug: true, - }, - Rule { - when: |mapping| mapping.error_str.contains("A timeout occurred"), - kind: Kind::Timeout(None), - message: |mapping| prefixed(mapping, ""), - debug: true, - }, - Rule { - when: |mapping| { + PublicError::NotFound, + ), + Rule::new( + |mapping| mapping.error_str.contains("A timeout occurred"), + PublicError::Timeout { status: 408 }, + ), + Rule::new( + |mapping| { let error_str = &mapping.error_str; (error_str.contains("invalid_request_error") && error_str.contains("content_policy_violation")) @@ -89,38 +35,29 @@ const RULES: &[Rule] = &[ .to_lowercase() .contains("request was rejected as a result of the safety system") }, - kind: with_response(StatusClass::ContentPolicyViolation), - message: |mapping| prefixed(mapping, "ContentPolicyViolationError: "), - debug: true, - }, + PublicError::ContentPolicyViolation, + ), Rule { - when: |mapping| { - contains_any( - &mapping.error_str, - &["invalid_encrypted_content", "could not be verified"], - ) - }, - kind: with_response(StatusClass::BadRequest), - message: |mapping| { - format!( - "{} - {}{ENCRYPTED_CONTENT_HELP}", - exception_provider(mapping), - message(mapping) - ) - }, - debug: true, + hint: ENCRYPTED_CONTENT_HELP, + ..Rule::new( + |mapping| { + contains_any( + &mapping.error_str, + &["invalid_encrypted_content", "could not be verified"], + ) + }, + PublicError::BadRequest, + ) }, - Rule { - when: |mapping| { + Rule::new( + |mapping| { mapping.error_str.contains("invalid_request_error") && !mapping.error_str.contains("Incorrect API key provided") }, - kind: with_response(StatusClass::BadRequest), - message: |mapping| prefixed(mapping, ""), - debug: true, - }, - Rule { - when: |mapping| { + PublicError::BadRequest, + ), + Rule::new( + |mapping| { contains_any( &mapping.error_str, &[ @@ -129,458 +66,127 @@ const RULES: &[Rule] = &[ ], ) }, - kind: Kind::Status { - class: StatusClass::InternalServer, - response: ResponseChoice::Omitted, - }, - message: |mapping| prefixed(mapping, ""), - debug: false, - }, - Rule { - when: |mapping| mapping.error_str.contains("Request too large"), - kind: with_response(StatusClass::RateLimit), - message: |mapping| prefixed(mapping, "RateLimitError: "), - debug: true, - }, - Rule { - when: |mapping| { - mapping.error_str.contains("The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable") - }, - kind: with_response(StatusClass::Authentication), - message: |mapping| prefixed(mapping, "AuthenticationError: "), - debug: true, - }, - Rule { - when: |mapping| { + PublicError::InternalServer, + ), + Rule::new( + |mapping| mapping.error_str.contains("Request too large"), + PublicError::RateLimit, + ), + Rule::new( + |mapping| { mapping .error_str .contains("Mistral API raised a streaming error") }, - kind: Kind::Api { - status: ApiStatus::Fixed(500), - request_url: OPENAI_URL, - }, - message: |mapping| prefixed(mapping, ""), - debug: true, - }, - Rule { - when: |mapping| mapping.original.status.is_none(), - kind: Kind::ApiConnection, - message: |mapping| prefixed(mapping, "APIConnectionError: "), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[400, 422]), - kind: with_response(StatusClass::BadRequest), - message: |mapping| prefixed(mapping, ""), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[401]), - kind: with_response(StatusClass::Authentication), - message: |mapping| prefixed(mapping, "AuthenticationError: "), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[404]), - kind: with_response(StatusClass::NotFound), - message: |mapping| prefixed(mapping, "NotFoundError: "), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[408]), - kind: Kind::Timeout(None), - message: |mapping| prefixed(mapping, "Timeout Error: "), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[429]), - kind: with_response(StatusClass::RateLimit), - message: |mapping| prefixed(mapping, "RateLimitError: "), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[500]), - kind: with_response(StatusClass::InternalServer), - message: |mapping| prefixed(mapping, "InternalServerError: "), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[502]), - kind: with_response(StatusClass::BadGateway), - message: |mapping| prefixed(mapping, "BadGatewayError: "), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[503]), - kind: with_response(StatusClass::ServiceUnavailable), - message: |mapping| prefixed(mapping, "ServiceUnavailableError: "), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, &[504]), - kind: Kind::Timeout(Some(504)), - message: |mapping| prefixed(mapping, "Timeout Error: "), - debug: true, - }, - Rule { - when: |_| true, - kind: Kind::Api { - status: ApiStatus::Original, - request_url: DOCS_URL, - }, - message: |mapping| prefixed(mapping, "APIError: "), - debug: true, - }, + PublicError::Api { status: 500 }, + ), ]; -pub(super) fn map(mapping: &Mapping<'_>) -> Option { - apply(RULES, mapping) -} - #[cfg(test)] mod tests { - use super::super::testing::{context, failure, http, status, upstream, with_debug}; - use super::super::{ExceptionFamily, OriginalException, PublicKind}; + use super::super::rules::first_match; + use super::super::testing::mapping; use super::*; - fn mapped(provider: &str, original: &OriginalException) -> PublicFailure { - let context = context(provider, ExceptionFamily::OpenAiCompatible); - map(&Mapping::new(&context, original)).expect("the OpenAI table ends in a catch-all") - } - - fn kind(class: StatusClass, status_code: u16, body: &str) -> PublicKind { - status(class, upstream(status_code, body)) + fn classified(status: Option, text: &str) -> Option { + first_match(RULES, &mapping(status, text)).map(|rule| rule.error) } #[rstest::rstest] - #[case::rate_limit_phrase( - 400, - "rate limit reached", - failure( - kind(StatusClass::RateLimit, 400, "rate limit reached"), - "RateLimitError: MistralException - rate limit reached", - "mistral", - ) - )] + #[case::rate_limit_phrase("rate limit reached", PublicError::RateLimit)] #[case::context_window( - 500, "This model's maximum context length is 10", - with_debug(failure( - kind( - StatusClass::ContextWindowExceeded, - 500, - "This model's maximum context length is 10" - ), - "ContextWindowExceededError: MistralException - This model's maximum context length is 10", - "mistral", - )) + PublicError::ContextWindowExceeded )] - #[case::model_not_found( - 400, - "invalid_request_error model_not_found", - with_debug(failure( - kind(StatusClass::NotFound, 400, "invalid_request_error model_not_found"), - "MistralException - invalid_request_error model_not_found", - "mistral", - )) - )] - #[case::timeout_occurred(400, "A timeout occurred", with_debug(failure( - PublicKind::Timeout { status: None }, - "MistralException - A timeout occurred", - "mistral", - )))] + #[case::model_not_found("invalid_request_error model_not_found", PublicError::NotFound)] + #[case::timeout_occurred("A timeout occurred", PublicError::Timeout { status: 408 })] #[case::content_policy_error_code( - 400, "invalid_request_error content_policy_violation", - with_debug(failure( - kind( - StatusClass::ContentPolicyViolation, - 400, - "invalid_request_error content_policy_violation" - ), - "ContentPolicyViolationError: MistralException - invalid_request_error content_policy_violation", - "mistral", - )) + PublicError::ContentPolicyViolation )] #[case::content_policy_usage_policy( - 400, "Invalid prompt violating our usage policy", - with_debug(failure( - kind( - StatusClass::ContentPolicyViolation, - 400, - "Invalid prompt violating our usage policy" - ), - "ContentPolicyViolationError: MistralException - Invalid prompt violating our usage policy", - "mistral", - )) + PublicError::ContentPolicyViolation )] #[case::content_policy_safety_system( - 400, "Request was rejected as a result of the safety system", - with_debug(failure( - kind( - StatusClass::ContentPolicyViolation, - 400, - "Request was rejected as a result of the safety system" - ), - "ContentPolicyViolationError: MistralException - Request was rejected as a result of the safety system", - "mistral", - )) - )] - #[case::encrypted_content(400, "invalid_encrypted_content", with_debug(failure( - kind(StatusClass::BadRequest, 400, "invalid_encrypted_content"), - &format!("MistralException - invalid_encrypted_content{ENCRYPTED_CONTENT_HELP}"), - "mistral", - )))] - #[case::unverifiable_content(400, "could not be verified", with_debug(failure( - kind(StatusClass::BadRequest, 400, "could not be verified"), - &format!("MistralException - could not be verified{ENCRYPTED_CONTENT_HELP}"), - "mistral", - )))] - #[case::invalid_request( - 429, - "invalid_request_error bad field", - with_debug(failure( - kind(StatusClass::BadRequest, 429, "invalid_request_error bad field"), - "MistralException - invalid_request_error bad field", - "mistral", - )) + PublicError::ContentPolicyViolation )] + #[case::encrypted_content("invalid_encrypted_content", PublicError::BadRequest)] + #[case::unverifiable_content("could not be verified", PublicError::BadRequest)] + #[case::invalid_request("invalid_request_error bad field", PublicError::BadRequest)] #[case::unknown_server_error( - 400, "Web server is returning an unknown error", - failure( - status(StatusClass::InternalServer, None), - "MistralException - Web server is returning an unknown error", - "mistral", - ) + PublicError::InternalServer )] #[case::server_had_an_error( - 400, "The server had an error processing your request.", - failure( - status(StatusClass::InternalServer, None), - "MistralException - The server had an error processing your request.", - "mistral", - ) + PublicError::InternalServer )] - #[case::request_too_large( - 400, - "Request too large", - with_debug(failure( - kind(StatusClass::RateLimit, 400, "Request too large"), - "RateLimitError: MistralException - Request too large", - "mistral", - )) + #[case::request_too_large("Request too large", PublicError::RateLimit)] + #[case::mistral_streaming_error( + "Mistral API raised a streaming error", + PublicError::Api { status: 500 } )] - #[case::missing_client_api_key( - 400, - "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable", - with_debug(failure( - kind( - StatusClass::Authentication, - 400, - "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable", - ), - "AuthenticationError: MistralException - The api_key client option must be set either by passing api_key to the client or by setting the MISTRAL_API_KEY environment variable", - "mistral", - )) - )] - #[case::mistral_streaming_error(400, "Mistral API raised a streaming error", with_debug(failure( - PublicKind::Api { status: 500, request_url: OPENAI_URL }, - "MistralException - Mistral API raised a streaming error", - "mistral", - )))] - fn each_text_rule_maps_by_the_body( - #[case] status_code: u16, - #[case] body: &str, - #[case] expected: PublicFailure, - ) { - assert_eq!(mapped("mistral", &http(status_code, body)), expected); - } - - #[rstest::rstest] - #[case::bad_request( - 400, - kind(StatusClass::BadRequest, 400, "rejected"), - "MistralException - rejected" - )] - #[case::unprocessable( - 422, - kind(StatusClass::BadRequest, 422, "rejected"), - "MistralException - rejected" - )] - #[case::authentication( - 401, - kind(StatusClass::Authentication, 401, "rejected"), - "AuthenticationError: MistralException - rejected" - )] - #[case::not_found( - 404, - kind(StatusClass::NotFound, 404, "rejected"), - "NotFoundError: MistralException - rejected" - )] - #[case::request_timeout(408, PublicKind::Timeout { status: None }, "Timeout Error: MistralException - rejected")] - #[case::rate_limited( - 429, - kind(StatusClass::RateLimit, 429, "rejected"), - "RateLimitError: MistralException - rejected" - )] - #[case::internal_server( - 500, - kind(StatusClass::InternalServer, 500, "rejected"), - "InternalServerError: MistralException - rejected" - )] - #[case::bad_gateway( - 502, - kind(StatusClass::BadGateway, 502, "rejected"), - "BadGatewayError: MistralException - rejected" - )] - #[case::service_unavailable( - 503, - kind(StatusClass::ServiceUnavailable, 503, "rejected"), - "ServiceUnavailableError: MistralException - rejected" - )] - #[case::gateway_timeout(504, PublicKind::Timeout { status: Some(504) }, "Timeout Error: MistralException - rejected")] - #[case::any_other_status(409, PublicKind::Api { status: 409, request_url: DOCS_URL }, "APIError: MistralException - rejected")] - fn each_status_rule_maps_by_the_status( - #[case] status_code: u16, - #[case] kind: PublicKind, - #[case] message: &str, - ) { - assert_eq!( - mapped("mistral", &http(status_code, "rejected")), - with_debug(failure(kind, message, "mistral")) - ); - } - - #[test] - fn a_failure_without_a_status_is_a_connection_error() { - let original = OriginalException::Response { - message: "invalid OCR response field: pages".into(), - }; - assert_eq!( - mapped("mistral", &original), - with_debug(failure( - PublicKind::ApiConnection, - "APIConnectionError: MistralException - invalid OCR response field: pages", - "mistral" - )) - ); + fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) { + assert_eq!(classified(Some(400), text), Some(expected)); } #[rstest::rstest] #[case::rate_limit_before_context_window( - 400, "rate limit and This model's maximum context length is 10", - kind( - StatusClass::RateLimit, - 400, - "rate limit and This model's maximum context length is 10" - ), - "RateLimitError: MistralException - rate limit and This model's maximum context length is 10", - false + PublicError::RateLimit )] #[case::context_window_before_content_policy( - 400, "This model's maximum context length is 10 invalid_request_error content_policy_violation", - kind( - StatusClass::ContextWindowExceeded, - 400, - "This model's maximum context length is 10 invalid_request_error content_policy_violation" - ), - "ContextWindowExceededError: MistralException - This model's maximum context length is 10 invalid_request_error content_policy_violation", - true + PublicError::ContextWindowExceeded )] #[case::model_not_found_before_invalid_request( - 400, "invalid_request_error model_not_found", - kind(StatusClass::NotFound, 400, "invalid_request_error model_not_found"), - "MistralException - invalid_request_error model_not_found", - true + PublicError::NotFound )] #[case::timeout_before_invalid_request( - 400, "A timeout occurred invalid_request_error", - PublicKind::Timeout { status: None }, - "MistralException - A timeout occurred invalid_request_error", - true + PublicError::Timeout { status: 408 } )] #[case::content_policy_before_invalid_request( - 400, "invalid_request_error content_policy_violation", - kind( - StatusClass::ContentPolicyViolation, - 400, - "invalid_request_error content_policy_violation" - ), - "ContentPolicyViolationError: MistralException - invalid_request_error content_policy_violation", - true + PublicError::ContentPolicyViolation )] - #[case::invalid_request_with_a_bad_key_falls_to_the_status( - 401, - "invalid_request_error Incorrect API key provided", - kind( - StatusClass::Authentication, - 401, - "invalid_request_error Incorrect API key provided" - ), - "AuthenticationError: MistralException - invalid_request_error Incorrect API key provided", - true + #[case::encrypted_content_before_invalid_request( + "invalid_request_error invalid_encrypted_content", + PublicError::BadRequest )] - #[case::text_rules_before_status( - 429, - "invalid_request_error bad field", - kind(StatusClass::BadRequest, 429, "invalid_request_error bad field"), - "MistralException - invalid_request_error bad field", - true - )] - #[case::echoed_429_is_not_a_rate_limit( - 400, - "token 429 in the prompt", - kind(StatusClass::BadRequest, 400, "token 429 in the prompt"), - "MistralException - token 429 in the prompt", - true - )] - fn the_earlier_rule_wins_when_two_apply( - #[case] status_code: u16, - #[case] body: &str, - #[case] kind: PublicKind, - #[case] message: &str, - #[case] debug: bool, + fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) { + assert_eq!(classified(Some(400), text), Some(expected)); + } + + #[rstest::rstest] + #[case::encrypted_content("invalid_encrypted_content", ENCRYPTED_CONTENT_HELP)] + #[case::plain_invalid_request("invalid_request_error bad field", "")] + fn only_encrypted_content_failures_carry_the_affinity_help( + #[case] text: &str, + #[case] hint: &str, ) { - let expected = failure(kind, message, "mistral"); assert_eq!( - mapped("mistral", &http(status_code, body)), - if debug { - with_debug(expected) - } else { - expected - } + first_match(RULES, &mapping(Some(400), text)).map(|rule| rule.hint), + Some(hint) ); } #[rstest::rstest] - #[case::provider_names_replace_openai( - "azure_ai", - "OPENAI said openai.OpenAIError", - "Azure_aiException - AZURE_AI said azure_ai.azure_aiError" - )] - #[case::openai_keeps_its_own_name("openai", "rejected", "OpenAIException - rejected")] - fn the_message_names_the_provider( - #[case] provider: &str, - #[case] body: &str, - #[case] message: &str, - ) { + #[case::bad_key_is_left_to_the_status("invalid_request_error Incorrect API key provided")] + #[case::echoed_429_is_not_a_rate_limit("token 429 in the prompt")] + #[case::unmarked("rejected")] + fn text_without_a_marker_is_left_to_the_status_table(#[case] text: &str) { + assert_eq!(classified(Some(400), text), None); + } + + #[test] + fn a_standalone_429_counts_only_with_a_429_status() { assert_eq!( - mapped(provider, &http(400, body)), - with_debug(failure( - kind(StatusClass::BadRequest, 400, body), - message, - provider - )) + classified(Some(429), "got 429 back"), + Some(PublicError::RateLimit) ); } } diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs index d64868c1606..82392cbd7ee 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs @@ -1,14 +1,4 @@ -use super::public::StatusClass; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum LocalClass { - ValueError, - FileNotFound, - OsError, -} - -/// A route failure in the shape Python's `exception_type` receives it, before any public -/// class is chosen. +/// A failure a Rust route produced, before any public class is chosen. #[derive(Clone, Debug, PartialEq)] pub enum OriginalException { Http { @@ -23,27 +13,132 @@ pub enum OriginalException { timeout_seconds: Option, elapsed_seconds: Option, }, - Response { - message: String, - }, - Local { - class: LocalClass, - message: String, - }, - /// A failure Python raises as a public LiteLLM exception itself, which `exception_type` - /// hands back unchanged. - Public { - class: StatusClass, + /// A failure with no HTTP response behind it, such as an unparseable body or a local + /// file error. + Plain { message: String, }, } -/// Which of the provider-specific mappers in `exception_type` a route's provider uses. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +/// Which provider-specific text rules apply before the shared status table. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ExceptionFamily { OpenAiCompatible, VertexAi, Cohere, - #[default] Other, } + +/// `openai_compatible_providers` in `litellm/constants.py`. +const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[ + "anyscale", + "groq", + "nvidia_nim", + "cerebras", + "baseten", + "sambanova", + "ai21_chat", + "ai21", + "volcengine", + "codestral", + "deepseek", + "tencent", + "deepinfra", + "perplexity", + "xinference", + "xai", + "zai", + "together_ai", + "fireworks_ai", + "empower", + "friendliai", + "azure_ai", + "github", + "litellm_proxy", + "hosted_vllm", + "llamafile", + "lm_studio", + "galadriel", + "github_copilot", + "chatgpt", + "novita", + "meta_llama", + "publicai", + "synthetic", + "tensormesh", + "apertis", + "nano-gpt", + "poe", + "chutes", + "parasail", + "libertai", + "featherless_ai", + "nscale", + "nebius", + "dashscope", + "qwencloud", + "qwen_ai_platform", + "modelscope", + "moonshot", + "v0", + "helicone", + "morph", + "lambda_ai", + "inception", + "hyperbolic", + "vercel_ai_gateway", + "aiml", + "wandb", + "cometapi", + "clarifai", + "docker_model_runner", + "ragflow", + "pinstripes", + "darkbloom", + "meta", + "cognition", + "scx-ai", +]; + +impl ExceptionFamily { + /// The provider dispatch at the top of Python's `exception_type`, in its order. + pub fn for_provider(provider: &str) -> Self { + match provider { + "openai" | "text-completion-openai" | "custom_openai" | "mistral" | "runwayml" => { + Self::OpenAiCompatible + } + provider if OPENAI_COMPATIBLE_PROVIDERS.contains(&provider) => Self::OpenAiCompatible, + "vertex_ai" | "vertex_ai_beta" | "gemini" => Self::VertexAi, + "cohere" | "cohere_chat" => Self::Cohere, + _ => Self::Other, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[rstest::rstest] + #[case::openai("openai", ExceptionFamily::OpenAiCompatible)] + #[case::text_completion_openai("text-completion-openai", ExceptionFamily::OpenAiCompatible)] + #[case::custom_openai("custom_openai", ExceptionFamily::OpenAiCompatible)] + #[case::mistral("mistral", ExceptionFamily::OpenAiCompatible)] + #[case::runwayml("runwayml", ExceptionFamily::OpenAiCompatible)] + #[case::listed_compatible("azure_ai", ExceptionFamily::OpenAiCompatible)] + #[case::compatible_list_wins_over_its_own_mapper( + "together_ai", + ExceptionFamily::OpenAiCompatible + )] + #[case::vertex_ai("vertex_ai", ExceptionFamily::VertexAi)] + #[case::vertex_ai_beta("vertex_ai_beta", ExceptionFamily::VertexAi)] + #[case::gemini("gemini", ExceptionFamily::VertexAi)] + #[case::cohere("cohere", ExceptionFamily::Cohere)] + #[case::cohere_chat("cohere_chat", ExceptionFamily::Cohere)] + #[case::unported_mapper("anthropic", ExceptionFamily::Other)] + #[case::unknown("reducto", ExceptionFamily::Other)] + #[case::empty("", ExceptionFamily::Other)] + fn provider_selects_the_family(#[case] provider: &str, #[case] family: ExceptionFamily) { + assert_eq!(ExceptionFamily::for_provider(provider), family); + } +} diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/public.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/public.rs index a7319c1287b..c567185aa6d 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/public.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/public.rs @@ -1,262 +1,77 @@ -use serde::Serialize; - -/// The public LiteLLM classes built from a status code alone: every one takes the same -/// constructor arguments. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::EnumIter, strum::IntoStaticStr)] -#[serde(rename_all = "snake_case")] -#[strum(serialize_all = "snake_case")] -pub enum StatusClass { +/// The public LiteLLM exception classes a Rust route failure can become. Python builds the +/// class; Rust decides which one. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PublicError { BadRequest, + ContextWindowExceeded, + ContentPolicyViolation, Authentication, PermissionDenied, NotFound, + Timeout { status: u16 }, RateLimit, - ContextWindowExceeded, - ContentPolicyViolation, InternalServer, BadGateway, ServiceUnavailable, - UnsupportedParams, + ApiConnection, + Api { status: u16 }, } -impl StatusClass { - /// The `status_code` the Python class sets on itself. +impl PublicError { + /// The `status_code` the Python class carries. pub const fn status_code(self) -> u16 { match self { - Self::BadRequest - | Self::ContextWindowExceeded - | Self::ContentPolicyViolation - | Self::UnsupportedParams => 400, + Self::BadRequest | Self::ContextWindowExceeded | Self::ContentPolicyViolation => 400, Self::Authentication => 401, Self::PermissionDenied => 403, Self::NotFound => 404, Self::RateLimit => 429, - Self::InternalServer => 500, + Self::InternalServer | Self::ApiConnection => 500, Self::BadGateway => 502, Self::ServiceUnavailable => 503, + Self::Timeout { status } | Self::Api { status } => status, } } } -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct UpstreamResponse { pub status: u16, pub body: String, pub headers: Vec<(String, String)>, } -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -pub struct HttpStub { - pub status: u16, - pub method: &'static str, - pub url: &'static str, - pub content: Option, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum ResponseArg { - Upstream(UpstreamResponse), - Stub(HttpStub), -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum PublicKind { - Status { - status_class: StatusClass, - response: Option, - }, - Timeout { - status: Option, - }, - ApiConnection, - Api { - status: u16, - request_url: &'static str, - }, -} - -/// Constructor arguments for the public LiteLLM exception, as `exception_type` passes them. -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -pub struct PublicFailure { - pub kind: PublicKind, +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MappedFailure { + pub error: PublicError, pub message: String, - pub model: String, - pub llm_provider: Option, - pub litellm_debug_info: Option, - pub litellm_response_headers: Option>, - pub print_banner: bool, + pub upstream: Option, + pub debug_info: String, } #[cfg(test)] mod tests { - use std::collections::BTreeSet; - use std::path::PathBuf; - - use serde_json::Value; - use strum::IntoEnumIterator; - use super::*; - const REGENERATE: &str = "LITELLM_REGENERATE_PUBLIC_FAILURE_FIXTURES"; - - fn fixture_directory() -> PathBuf { - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("../../../tests/test_litellm/rust_bridge/fixtures/public_failures") - } - - fn upstream(status: u16) -> ResponseArg { - ResponseArg::Upstream(UpstreamResponse { - status, - body: r#"{"message": "rejected"}"#.into(), - headers: vec![("retry-after".into(), "7".into())], - }) - } - - fn status_response(class: StatusClass) -> Option { - match class { - StatusClass::Authentication => None, - StatusClass::PermissionDenied => Some(ResponseArg::Stub(HttpStub { - status: 403, - method: "POST", - url: " https://cloud.google.com/vertex-ai/", - content: None, - })), - StatusClass::InternalServer => Some(ResponseArg::Stub(HttpStub { - status: 500, - method: "completion", - url: "https://github.com/BerriAI/litellm", - content: Some("upstream text".into()), - })), - class => Some(upstream(class.status_code())), - } - } - - fn failure(kind: PublicKind, name: &str) -> PublicFailure { - let headers = matches!( - &kind, - PublicKind::Status { - response: Some(ResponseArg::Upstream(_)), - .. - } - ); - PublicFailure { - kind, - message: format!("MistralException - {name}"), - model: "ocr-model".into(), - llm_provider: Some("mistral".into()), - litellm_debug_info: Some("\nModel: ocr-model".into()), - litellm_response_headers: headers.then(|| vec![("retry-after".into(), "7".into())]), - print_banner: false, - } - } - - /// One payload per public class the constructor can build; `test_failures.py` reads the - /// same files, so a shape change on either side fails there or here. - fn fixtures() -> Vec<(String, PublicFailure)> { - let statuses = StatusClass::iter().map(|class| { - let name: &'static str = class.into(); - let name = format!("status_{name}"); - let built = failure( - PublicKind::Status { - status_class: class, - response: status_response(class), - }, - &name, - ); - (name, built) - }); - let others = [ - ( - "timeout_with_status", - PublicFailure { - print_banner: true, - ..failure( - PublicKind::Timeout { status: Some(504) }, - "timeout_with_status", - ) - }, - ), - ( - "timeout_without_status", - PublicFailure { - litellm_debug_info: None, - ..failure( - PublicKind::Timeout { status: None }, - "timeout_without_status", - ) - }, - ), - ( - "api_connection", - PublicFailure { - llm_provider: None, - ..failure(PublicKind::ApiConnection, "api_connection") - }, - ), - ( - "api", - failure( - PublicKind::Api { - status: 409, - request_url: "https://docs.litellm.ai/docs", - }, - "api", - ), - ), - ] - .map(|(name, built)| (name.to_string(), built)); - statuses.chain(others).collect() - } - - #[test] - fn serialized_payloads_match_the_golden_fixtures_python_reads() { - let directory = fixture_directory(); - let regenerate = std::env::var_os(REGENERATE).is_some(); - let expected = fixtures(); - for (name, built) in &expected { - let path = directory.join(format!("{name}.json")); - let serialized = serde_json::to_value(built).unwrap(); - if regenerate { - std::fs::create_dir_all(&directory).unwrap(); - std::fs::write( - &path, - format!("{}\n", serde_json::to_string_pretty(&serialized).unwrap()), - ) - .unwrap(); - } - let golden: Value = - serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); - assert_eq!(serialized, golden, "{name}; set {REGENERATE}=1 to rewrite"); - } - let on_disk: BTreeSet = std::fs::read_dir(&directory) - .unwrap() - .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned()) - .collect(); - let generated: BTreeSet = expected - .iter() - .map(|(name, _)| format!("{name}.json")) - .collect(); - assert_eq!(on_disk, generated); - } - #[rstest::rstest] - #[case(StatusClass::BadRequest, 400)] - #[case(StatusClass::Authentication, 401)] - #[case(StatusClass::PermissionDenied, 403)] - #[case(StatusClass::NotFound, 404)] - #[case(StatusClass::RateLimit, 429)] - #[case(StatusClass::ContextWindowExceeded, 400)] - #[case(StatusClass::ContentPolicyViolation, 400)] - #[case(StatusClass::InternalServer, 500)] - #[case(StatusClass::BadGateway, 502)] - #[case(StatusClass::ServiceUnavailable, 503)] - #[case(StatusClass::UnsupportedParams, 400)] + #[case::bad_request(PublicError::BadRequest, 400)] + #[case::context_window(PublicError::ContextWindowExceeded, 400)] + #[case::content_policy(PublicError::ContentPolicyViolation, 400)] + #[case::authentication(PublicError::Authentication, 401)] + #[case::permission_denied(PublicError::PermissionDenied, 403)] + #[case::not_found(PublicError::NotFound, 404)] + #[case::request_timeout(PublicError::Timeout { status: 408 }, 408)] + #[case::gateway_timeout(PublicError::Timeout { status: 504 }, 504)] + #[case::rate_limit(PublicError::RateLimit, 429)] + #[case::internal_server(PublicError::InternalServer, 500)] + #[case::api_connection(PublicError::ApiConnection, 500)] + #[case::bad_gateway(PublicError::BadGateway, 502)] + #[case::service_unavailable(PublicError::ServiceUnavailable, 503)] + #[case::api(PublicError::Api { status: 501 }, 501)] fn status_codes_are_the_ones_the_python_classes_set( - #[case] class: StatusClass, + #[case] error: PublicError, #[case] status: u16, ) { - assert_eq!(class.status_code(), status); + assert_eq!(error.status_code(), status); } } diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs index 34cf2c390c3..0346a8bc718 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/rules.rs @@ -4,106 +4,29 @@ use fancy_regex::Regex; use serde_json::Value; use super::Mapping; -use super::public::{HttpStub, PublicFailure, PublicKind, ResponseArg, StatusClass}; +use super::public::PublicError; -const GITHUB_URL: &str = "https://github.com/BerriAI/litellm"; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(super) enum ResponseChoice { - Omitted, - Provider, - Stub { status: u16, url: &'static str }, - InternalServerStub, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(super) enum ApiStatus { - Fixed(u16), - Original, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(super) enum Kind { - Status { - class: StatusClass, - response: ResponseChoice, - }, - Timeout(Option), - ApiConnection, - Api { - status: ApiStatus, - request_url: &'static str, - }, -} - -/// One branch of a Python `_map_*_exception` function: when it applies, the class it -/// raises, the message it builds, and whether it passes `litellm_debug_info`. +/// One text branch of a Python `_map_*_exception` function: when it applies, the class it +/// raises, and any help text appended to the message. pub(super) struct Rule { - pub(super) when: fn(&Mapping<'_>) -> bool, - pub(super) kind: Kind, - pub(super) message: fn(&Mapping<'_>) -> String, - pub(super) debug: bool, -} - -/// The first rule that applies decides the failure, as the `if`/`elif` chain does in Python. -pub(super) fn apply(rules: &[Rule], mapping: &Mapping<'_>) -> Option { - rules - .iter() - .find(|rule| (rule.when)(mapping)) - .map(|rule| rule.build(mapping)) + pub(super) when: fn(&Mapping) -> bool, + pub(super) error: PublicError, + pub(super) hint: &'static str, } impl Rule { - fn build(&self, mapping: &Mapping<'_>) -> PublicFailure { - let kind = match self.kind { - Kind::Status { class, response } => PublicKind::Status { - status_class: class, - response: response.resolve(mapping), - }, - Kind::Timeout(status) => PublicKind::Timeout { status }, - Kind::ApiConnection => PublicKind::ApiConnection, - Kind::Api { - status, - request_url, - } => PublicKind::Api { - status: match status { - ApiStatus::Fixed(status) => status, - ApiStatus::Original => mapping.original.status.unwrap_or(500), - }, - request_url, - }, - }; - PublicFailure { - kind, - message: (self.message)(mapping), - model: mapping.context.model.clone(), - llm_provider: mapping.context.custom_llm_provider.clone(), - litellm_debug_info: self.debug.then(|| mapping.extra_information.clone()), - litellm_response_headers: None, - print_banner: false, + pub(super) const fn new(when: fn(&Mapping) -> bool, error: PublicError) -> Self { + Self { + when, + error, + hint: "", } } } -impl ResponseChoice { - fn resolve(self, mapping: &Mapping<'_>) -> Option { - match self { - Self::Omitted => None, - Self::Provider => mapping.original.response.clone().map(ResponseArg::Upstream), - Self::Stub { status, url } => Some(ResponseArg::Stub(HttpStub { - status, - method: "POST", - url, - content: None, - })), - Self::InternalServerStub => Some(ResponseArg::Stub(HttpStub { - status: 500, - method: "completion", - url: GITHUB_URL, - content: Some(mapping.original.message.clone()), - })), - } - } +/// The first rule that applies decides the class, as the `if`/`elif` chain does in Python. +pub(super) fn first_match<'r>(rules: &'r [Rule], mapping: &Mapping) -> Option<&'r Rule> { + rules.iter().find(|rule| (rule.when)(mapping)) } pub(super) fn contains_any(text: &str, markers: &[&str]) -> bool { @@ -117,7 +40,7 @@ static RATE_LIMIT_PHRASE: LazyLock = /// `ExceptionCheckers.is_error_str_rate_limit`. pub(super) fn is_rate_limit(error_str: &str, status: Option) -> bool { - if STANDALONE_429.is_match(error_str).unwrap_or(false) && status == Some(429) { + if STANDALONE_429.is_match(error_str).unwrap_or(false) && matches!(status, None | Some(429)) { return true; } let lower = error_str.to_lowercase(); @@ -169,184 +92,34 @@ pub(super) fn body_error_code(error_str: &str) -> Option { #[cfg(test)] mod tests { - use super::super::testing::{context, failure, http}; - use super::super::{ExceptionFamily, OriginalException, UpstreamResponse}; + use super::super::testing::mapping; use super::*; - fn first_marker(mapping: &Mapping<'_>) -> bool { - mapping.error_str.contains("first") - } - - fn always(_: &Mapping<'_>) -> bool { - true - } - - fn text(mapping: &Mapping<'_>) -> String { - format!("seen {}", mapping.error_str) - } - const ORDERED: &[Rule] = &[ - Rule { - when: first_marker, - kind: Kind::Status { - class: StatusClass::NotFound, - response: ResponseChoice::Omitted, - }, - message: text, - debug: false, - }, - Rule { - when: always, - kind: Kind::ApiConnection, - message: text, - debug: true, - }, + Rule::new( + |mapping| mapping.error_str.contains("first"), + PublicError::NotFound, + ), + Rule::new(|_| true, PublicError::ApiConnection), ]; - fn apply_one(kind: Kind, debug: bool, original: &OriginalException) -> Option { - let context = context("mistral", ExceptionFamily::OpenAiCompatible); - let mapping = Mapping::new(&context, original); - apply( - &[Rule { - when: always, - kind, - message: text, - debug, - }], - &mapping, - ) - } - #[rstest::rstest] - #[case::earlier_rule_wins("first and second", failure( - PublicKind::Status { status_class: StatusClass::NotFound, response: None }, - "seen first and second", - "mistral", - ))] - #[case::later_rule_when_the_earlier_does_not_apply("second", PublicFailure { - litellm_debug_info: Some("\nModel: ocr-model".into()), - ..failure(PublicKind::ApiConnection, "seen second", "mistral") - })] - fn the_first_applicable_rule_decides(#[case] body: &str, #[case] expected: PublicFailure) { - let context = context("mistral", ExceptionFamily::OpenAiCompatible); - let original = http(400, body); - assert_eq!( - apply(ORDERED, &Mapping::new(&context, &original)), - Some(expected) - ); + #[case::earlier_rule_wins("first and second", PublicError::NotFound)] + #[case::later_rule_when_the_earlier_does_not_apply("second", PublicError::ApiConnection)] + fn the_first_applicable_rule_decides(#[case] text: &str, #[case] expected: PublicError) { + let rule = first_match(ORDERED, &mapping(Some(400), text)); + assert_eq!(rule.map(|rule| rule.error), Some(expected)); } #[test] fn no_applicable_rule_leaves_the_failure_to_the_caller() { - let context = context("mistral", ExceptionFamily::OpenAiCompatible); - let original = http(400, "second"); - assert_eq!( - apply(&ORDERED[..1], &Mapping::new(&context, &original)), - None - ); - } - - #[rstest::rstest] - #[case::omitted(ResponseChoice::Omitted, None)] - #[case::provider(ResponseChoice::Provider, Some(ResponseArg::Upstream(UpstreamResponse { - status: 400, - body: "body".into(), - headers: vec![("retry-after".into(), "7".into())], - })))] - #[case::stub( - ResponseChoice::Stub { status: 429, url: "https://stub.test" }, - Some(ResponseArg::Stub(HttpStub { status: 429, method: "POST", url: "https://stub.test", content: None })) - )] - #[case::internal_server_stub( - ResponseChoice::InternalServerStub, - Some(ResponseArg::Stub(HttpStub { - status: 500, - method: "completion", - url: GITHUB_URL, - content: Some("body".into()), - })) - )] - fn response_choices_resolve_against_the_original( - #[case] response: ResponseChoice, - #[case] expected: Option, - ) { - let built = apply_one( - Kind::Status { - class: StatusClass::BadRequest, - response, - }, - false, - &http(400, "body"), - ) - .unwrap(); - assert_eq!( - built.kind, - PublicKind::Status { - status_class: StatusClass::BadRequest, - response: expected, - } - ); - } - - #[rstest::rstest] - #[case::fixed(ApiStatus::Fixed(500), http(409, "body"), 500)] - #[case::original(ApiStatus::Original, http(409, "body"), 409)] - #[case::original_without_a_status( - ApiStatus::Original, - OriginalException::Response { message: "body".into() }, - 500 - )] - fn api_status_is_fixed_or_the_originals( - #[case] status: ApiStatus, - #[case] original: OriginalException, - #[case] expected: u16, - ) { - let built = apply_one( - Kind::Api { - status, - request_url: "https://api.test", - }, - false, - &original, - ) - .unwrap(); - assert_eq!( - built, - failure( - PublicKind::Api { - status: expected, - request_url: "https://api.test" - }, - "seen body", - "mistral" - ) - ); - } - - #[rstest::rstest] - #[case::with_debug(true, Some("\nModel: ocr-model"))] - #[case::without_debug(false, None)] - fn debug_rules_carry_the_extra_information( - #[case] debug: bool, - #[case] expected: Option<&str>, - ) { - let built = apply_one(Kind::Timeout(Some(504)), debug, &http(504, "body")).unwrap(); - assert_eq!( - built, - PublicFailure { - litellm_debug_info: expected.map(str::to_string), - ..failure( - PublicKind::Timeout { status: Some(504) }, - "seen body", - "mistral" - ) - } - ); + assert!(first_match(&ORDERED[..1], &mapping(Some(400), "second")).is_none()); } #[rstest::rstest] #[case::standalone_429_with_429_status("got 429 back", Some(429), true)] #[case::standalone_429_with_other_status("got 429 back", Some(400), false)] + #[case::standalone_429_with_unknown_status("got 429 back", None, true)] #[case::embedded_429("token4290", Some(429), false)] #[case::phrase_spaced("Rate Limit reached", None, true)] #[case::phrase_underscored("rate_limit", None, true)] diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/status.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/status.rs index 3d817fe9ae2..cb8924d6c2a 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/status.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/status.rs @@ -1,153 +1,49 @@ -use super::public::{PublicFailure, StatusClass}; -use super::rules::{ApiStatus, Kind, ResponseChoice, Rule, apply}; -use super::{DOCS_URL, Mapping}; +use super::public::PublicError; -const fn with_response(class: StatusClass) -> Kind { - Kind::Status { - class, - response: ResponseChoice::Provider, - } -} - -fn message(mapping: &Mapping<'_>) -> String { - format!("{} - {}", mapping.exception_provider, mapping.error_str) -} - -fn status(mapping: &Mapping<'_>) -> u16 { - mapping.original.status.unwrap_or_default() -} - -/// `_map_exception_by_status`, the fallback for a provider error no provider mapper claimed. -const RULES: &[Rule] = &[ - Rule { - when: |mapping| status(mapping) == 401, - kind: with_response(StatusClass::Authentication), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) == 403, - kind: with_response(StatusClass::PermissionDenied), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) == 404, - kind: with_response(StatusClass::NotFound), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) == 408, - kind: Kind::Timeout(None), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) == 429, - kind: with_response(StatusClass::RateLimit), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) == 500, - kind: with_response(StatusClass::InternalServer), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) == 502, - kind: with_response(StatusClass::BadGateway), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) == 503, - kind: with_response(StatusClass::ServiceUnavailable), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) == 504, - kind: Kind::Timeout(Some(504)), - message, - debug: true, - }, - Rule { - when: |mapping| status(mapping) < 500, - kind: with_response(StatusClass::BadRequest), - message, - debug: true, - }, - Rule { - when: |_| true, - kind: Kind::Api { - status: ApiStatus::Original, - request_url: DOCS_URL, - }, - message, - debug: true, - }, -]; - -/// Only a real provider status of 400 or more reaches the table; a status the HTTP handler -/// synthesized for a failure without a response does not. -pub(super) fn map(mapping: &Mapping<'_>) -> Option { - let status = mapping.original.status?; - if status < 400 || mapping.original.status_is_synthesized { - return None; - } - apply(RULES, mapping) +/// `_map_exception_by_status`, the one place a provider status picks a class. Statuses +/// below 400 are not failures the table claims. +pub(super) fn classify(status: u16) -> Option { + let error = match status { + ..400 => return None, + 401 => PublicError::Authentication, + 403 => PublicError::PermissionDenied, + 404 => PublicError::NotFound, + 408 | 504 => PublicError::Timeout { status }, + 429 => PublicError::RateLimit, + 500 => PublicError::InternalServer, + 502 => PublicError::BadGateway, + 503 => PublicError::ServiceUnavailable, + 400..500 => PublicError::BadRequest, + _ => PublicError::Api { status }, + }; + Some(error) } #[cfg(test)] mod tests { - use super::super::testing::{context, failure, http, upstream, with_debug}; - use super::super::{ExceptionFamily, OriginalException, PublicKind}; use super::*; - fn mapped(original: &OriginalException) -> Option { - let context = context("reducto", ExceptionFamily::Other); - map(&Mapping::new(&context, original)) - } - - fn classified(class: StatusClass, status_code: u16) -> PublicKind { - PublicKind::Status { - status_class: class, - response: upstream(status_code, "rejected"), - } - } - #[rstest::rstest] - #[case::authentication(401, classified(StatusClass::Authentication, 401))] - #[case::permission_denied(403, classified(StatusClass::PermissionDenied, 403))] - #[case::not_found(404, classified(StatusClass::NotFound, 404))] - #[case::request_timeout(408, PublicKind::Timeout { status: None })] - #[case::rate_limited(429, classified(StatusClass::RateLimit, 429))] - #[case::internal_server(500, classified(StatusClass::InternalServer, 500))] - #[case::bad_gateway(502, classified(StatusClass::BadGateway, 502))] - #[case::service_unavailable(503, classified(StatusClass::ServiceUnavailable, 503))] - #[case::gateway_timeout(504, PublicKind::Timeout { status: Some(504) })] - #[case::lowest_client_error(400, classified(StatusClass::BadRequest, 400))] - #[case::other_client_error(409, classified(StatusClass::BadRequest, 409))] - #[case::highest_client_error(499, classified(StatusClass::BadRequest, 499))] - #[case::other_server_error(501, PublicKind::Api { status: 501, request_url: DOCS_URL })] - fn every_mapped_status_and_the_fallback(#[case] status_code: u16, #[case] kind: PublicKind) { - assert_eq!( - mapped(&http(status_code, "rejected")), - Some(with_debug(failure( - kind, - "ReductoException - rejected", - "reducto" - ))) - ); - } - - #[rstest::rstest] - #[case::below_client_errors(http(399, "rejected"))] - #[case::synthesized(OriginalException::Connection { message: "refused".into() })] - #[case::no_status(OriginalException::Response { message: "bad body".into() })] - fn failures_the_table_does_not_claim(#[case] original: OriginalException) { - assert_eq!(mapped(&original), None); + #[case::below_client_errors(399, None)] + #[case::lowest_client_error(400, Some(PublicError::BadRequest))] + #[case::authentication(401, Some(PublicError::Authentication))] + #[case::permission_denied(403, Some(PublicError::PermissionDenied))] + #[case::not_found(404, Some(PublicError::NotFound))] + #[case::request_timeout(408, Some(PublicError::Timeout { status: 408 }))] + #[case::other_client_error(409, Some(PublicError::BadRequest))] + #[case::unprocessable(422, Some(PublicError::BadRequest))] + #[case::rate_limited(429, Some(PublicError::RateLimit))] + #[case::highest_client_error(499, Some(PublicError::BadRequest))] + #[case::internal_server(500, Some(PublicError::InternalServer))] + #[case::other_server_error(501, Some(PublicError::Api { status: 501 }))] + #[case::bad_gateway(502, Some(PublicError::BadGateway))] + #[case::service_unavailable(503, Some(PublicError::ServiceUnavailable))] + #[case::gateway_timeout(504, Some(PublicError::Timeout { status: 504 }))] + #[case::highest_server_error(599, Some(PublicError::Api { status: 599 }))] + fn every_mapped_status_and_the_fallback( + #[case] status: u16, + #[case] expected: Option, + ) { + assert_eq!(classify(status), expected); } } diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/vertex_ai.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/vertex_ai.rs index 0fa8c19c4f6..dab1adb2329 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/vertex_ai.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/vertex_ai.rs @@ -1,60 +1,17 @@ -use super::public::{PublicFailure, PublicKind, ResponseArg, StatusClass, UpstreamResponse}; -use super::rules::{ - Kind, ResponseChoice, Rule, apply, body_error_code, contains_any, is_context_window_exceeded, -}; -use super::{Mapping, python_capitalize}; - -const VERTEX_URL: &str = "https://cloud.google.com/vertex-ai/"; -const VERTEX_URL_WITH_SPACE: &str = " https://cloud.google.com/vertex-ai/"; +use super::public::PublicError; +use super::rules::{Rule, body_error_code, contains_any, is_context_window_exceeded}; const QUOTA_MARKERS: &[&str] = &[ "429 Quota exceeded", "Quota exceeded for", "Resource exhausted", - "IndexError: list index out of range", "429 Unable to submit request because the service is temporarily out of capacity.", ]; -const fn stubbed(class: StatusClass, status: u16, url: &'static str) -> Kind { - Kind::Status { - class, - response: ResponseChoice::Stub { status, url }, - } -} - -const fn bare(class: StatusClass) -> Kind { - Kind::Status { - class, - response: ResponseChoice::Omitted, - } -} - -/// `{Provider}Exception{label} - {error_str}` with Python's `str.capitalize()`. -fn capitalized(mapping: &Mapping<'_>, label: &str) -> String { - format!( - "{}Exception{label} - {}", - python_capitalize(mapping.provider), - mapping.error_str - ) -} - -/// `litellm.{Class}: {provider}Exception - {error_str}` with the provider as given. -fn litellm_prefixed(mapping: &Mapping<'_>, class: &str) -> String { - format!( - "litellm.{class}: {}Exception - {}", - mapping.provider, mapping.error_str - ) -} - -fn status_is(mapping: &Mapping<'_>, status: u16) -> bool { - mapping.original.status == Some(status) -} - -/// `_map_vertex_exception`, in its branch order. A failure no rule claims falls through -/// to the status table. -const RULES: &[Rule] = &[ - Rule { - when: |mapping| { +/// The text branches of `_map_vertex_exception`, in its order. +pub(super) const RULES: &[Rule] = &[ + Rule::new( + |mapping| { contains_any( &mapping.error_str, &[ @@ -63,54 +20,32 @@ const RULES: &[Rule] = &[ ], ) }, - kind: stubbed(StatusClass::BadRequest, 400, VERTEX_URL_WITH_SPACE), - message: |mapping| litellm_prefixed(mapping, "BadRequestError"), - debug: true, - }, - Rule { - when: |mapping| { + PublicError::BadRequest, + ), + Rule::new( + |mapping| { mapping .error_str .contains("400 Request payload size exceeds") + || is_context_window_exceeded(&mapping.error_str) }, - kind: bare(StatusClass::ContextWindowExceeded), - message: |mapping| capitalized(mapping, ""), - debug: false, - }, - Rule { - when: |mapping| is_context_window_exceeded(&mapping.error_str), - kind: bare(StatusClass::ContextWindowExceeded), - message: |mapping| format!("ContextWindowExceededError: {}", capitalized(mapping, "")), - debug: true, - }, - Rule { - when: |mapping| { + PublicError::ContextWindowExceeded, + ), + Rule::new( + |mapping| { contains_any( &mapping.error_str, &["None Unknown Error.", "Content has no parts."], ) }, - kind: Kind::Status { - class: StatusClass::InternalServer, - response: ResponseChoice::InternalServerStub, - }, - message: |mapping| litellm_prefixed(mapping, "InternalServerError"), - debug: true, - }, - Rule { - when: |mapping| mapping.error_str.contains("API key not valid."), - kind: bare(StatusClass::Authentication), - message: |mapping| capitalized(mapping, ""), - debug: true, - }, - Rule { - when: |mapping| mapping.error_str.contains("403"), - kind: stubbed(StatusClass::BadRequest, 403, VERTEX_URL_WITH_SPACE), - message: |mapping| capitalized(mapping, " BadRequestError"), - debug: true, - }, - Rule { - when: |mapping| { + PublicError::InternalServer, + ), + Rule::new( + |mapping| mapping.error_str.contains("API key not valid."), + PublicError::Authentication, + ), + Rule::new( + |mapping| { contains_any( &mapping.error_str, &[ @@ -119,456 +54,124 @@ const RULES: &[Rule] = &[ ], ) }, - kind: stubbed( - StatusClass::ContentPolicyViolation, - 400, - VERTEX_URL_WITH_SPACE, - ), - message: |mapping| capitalized(mapping, " ContentPolicyViolationError"), - debug: true, - }, - Rule { - when: |mapping| { + PublicError::ContentPolicyViolation, + ), + Rule::new( + |mapping| { contains_any(&mapping.error_str, QUOTA_MARKERS) || (mapping - .original .status .is_some_and(|status| (500..600).contains(&status)) && body_error_code(&mapping.error_str) == Some(429)) }, - kind: stubbed(StatusClass::RateLimit, 429, VERTEX_URL_WITH_SPACE), - message: |mapping| litellm_prefixed(mapping, "RateLimitError"), - debug: true, - }, - Rule { - when: |mapping| { + PublicError::RateLimit, + ), + Rule::new( + |mapping| { contains_any( &mapping.error_str, &["500 Internal Server Error", "The model is overloaded."], ) }, - kind: bare(StatusClass::InternalServer), - message: |mapping| litellm_prefixed(mapping, "InternalServerError"), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, 400), - kind: stubbed(StatusClass::BadRequest, 400, VERTEX_URL), - message: |mapping| capitalized(mapping, " BadRequestError"), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, 401), - kind: bare(StatusClass::Authentication), - message: |mapping| capitalized(mapping, ""), - debug: false, - }, - Rule { - when: |mapping| status_is(mapping, 403), - kind: stubbed(StatusClass::PermissionDenied, 403, VERTEX_URL), - message: |mapping| capitalized(mapping, ""), - debug: false, - }, - Rule { - when: |mapping| status_is(mapping, 404), - kind: bare(StatusClass::NotFound), - message: |mapping| capitalized(mapping, ""), - debug: false, - }, - Rule { - when: |mapping| status_is(mapping, 408), - kind: Kind::Timeout(None), - message: |mapping| capitalized(mapping, ""), - debug: false, - }, - Rule { - when: |mapping| status_is(mapping, 429), - kind: stubbed(StatusClass::RateLimit, 429, VERTEX_URL_WITH_SPACE), - message: |mapping| format!("litellm.RateLimitError: {}", capitalized(mapping, "")), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, 500), - kind: Kind::Status { - class: StatusClass::InternalServer, - response: ResponseChoice::InternalServerStub, - }, - message: |mapping| capitalized(mapping, " InternalServerError"), - debug: true, - }, - Rule { - when: |mapping| status_is(mapping, 502), - kind: Kind::ApiConnection, - message: |mapping| capitalized(mapping, ""), - debug: false, - }, - Rule { - when: |mapping| status_is(mapping, 503), - kind: bare(StatusClass::ServiceUnavailable), - message: |mapping| capitalized(mapping, ""), - debug: false, - }, + PublicError::InternalServer, + ), ]; -pub(super) fn map(mapping: &Mapping<'_>) -> Option { - apply(RULES, mapping).map(|failure| keep_upstream_response(mapping, failure)) -} - -/// Deliberate divergence from `_map_vertex_exception`, which replaces the provider response -/// with a stub and so drops the upstream body and `retry-after`. The response keeps the -/// status the public class carries. -fn keep_upstream_response(mapping: &Mapping<'_>, failure: PublicFailure) -> PublicFailure { - let (PublicKind::Status { status_class, .. }, Some(upstream), false) = ( - &failure.kind, - &mapping.original.response, - mapping.original.status_is_synthesized, - ) else { - return failure; - }; - PublicFailure { - kind: PublicKind::Status { - status_class: *status_class, - response: Some(ResponseArg::Upstream(UpstreamResponse { - status: status_class.status_code(), - ..upstream.clone() - })), - }, - ..failure - } -} - #[cfg(test)] mod tests { - use super::super::testing::{context, failure, http, status, upstream, with_debug}; - use super::super::{ExceptionFamily, HttpStub, OriginalException}; + use super::super::rules::first_match; + use super::super::testing::mapping; use super::*; - fn mapped(original: &OriginalException) -> Option { - let context = context("vertex_ai", ExceptionFamily::VertexAi); - map(&Mapping::new(&context, original)) - } - - fn kept(class: StatusClass, body: &str) -> PublicKind { - status(class, upstream(class.status_code(), body)) + fn classified(status: Option, text: &str) -> Option { + first_match(RULES, &mapping(status, text)).map(|rule| rule.error) } #[rstest::rstest] #[case::api_not_enabled( - 400, "Vertex AI API has not been used in project x", - with_debug(failure( - kept( - StatusClass::BadRequest, - "Vertex AI API has not been used in project x" - ), - "litellm.BadRequestError: vertex_aiException - Vertex AI API has not been used in project x", - "vertex_ai", - )) - )] - #[case::project_not_found( - 400, - "Unable to find your project", - with_debug(failure( - kept(StatusClass::BadRequest, "Unable to find your project"), - "litellm.BadRequestError: vertex_aiException - Unable to find your project", - "vertex_ai", - )) + PublicError::BadRequest )] + #[case::project_not_found("Unable to find your project", PublicError::BadRequest)] #[case::payload_too_large( - 400, "400 Request payload size exceeds the limit", - failure( - kept( - StatusClass::ContextWindowExceeded, - "400 Request payload size exceeds the limit" - ), - "Vertex_aiException - 400 Request payload size exceeds the limit", - "vertex_ai", - ) + PublicError::ContextWindowExceeded )] #[case::context_window( - 500, "This model's maximum context length is 10", - with_debug(failure( - kept( - StatusClass::ContextWindowExceeded, - "This model's maximum context length is 10" - ), - "ContextWindowExceededError: Vertex_aiException - This model's maximum context length is 10", - "vertex_ai", - )) - )] - #[case::unknown_error( - 400, - "None Unknown Error.", - with_debug(failure( - kept(StatusClass::InternalServer, "None Unknown Error."), - "litellm.InternalServerError: vertex_aiException - None Unknown Error.", - "vertex_ai", - )) - )] - #[case::no_parts( - 400, - "Content has no parts.", - with_debug(failure( - kept(StatusClass::InternalServer, "Content has no parts."), - "litellm.InternalServerError: vertex_aiException - Content has no parts.", - "vertex_ai", - )) - )] - #[case::api_key_not_valid( - 400, - "API key not valid.", - with_debug(failure( - kept(StatusClass::Authentication, "API key not valid."), - "Vertex_aiException - API key not valid.", - "vertex_ai", - )) - )] - #[case::forbidden_text( - 400, - "got a 403", - with_debug(failure( - kept(StatusClass::BadRequest, "got a 403"), - "Vertex_aiException BadRequestError - got a 403", - "vertex_ai", - )) - )] - #[case::response_blocked( - 400, - "The response was blocked.", - with_debug(failure( - kept(StatusClass::ContentPolicyViolation, "The response was blocked."), - "Vertex_aiException ContentPolicyViolationError - The response was blocked.", - "vertex_ai", - )) + PublicError::ContextWindowExceeded )] + #[case::unknown_error("None Unknown Error.", PublicError::InternalServer)] + #[case::no_parts("Content has no parts.", PublicError::InternalServer)] + #[case::api_key_not_valid("API key not valid.", PublicError::Authentication)] + #[case::response_blocked("The response was blocked.", PublicError::ContentPolicyViolation)] #[case::output_blocked( - 400, "Output blocked by content filtering policy", - with_debug(failure( - kept( - StatusClass::ContentPolicyViolation, - "Output blocked by content filtering policy" - ), - "Vertex_aiException ContentPolicyViolationError - Output blocked by content filtering policy", - "vertex_ai", - )) + PublicError::ContentPolicyViolation )] - #[case::quota_marker( - 400, - "Quota exceeded for aiplatform", - with_debug(failure( - kept(StatusClass::RateLimit, "Quota exceeded for aiplatform"), - "litellm.RateLimitError: vertex_aiException - Quota exceeded for aiplatform", - "vertex_ai", - )) + #[case::quota_exceeded_429("429 Quota exceeded", PublicError::RateLimit)] + #[case::quota_exceeded_for("Quota exceeded for aiplatform", PublicError::RateLimit)] + #[case::resource_exhausted("Resource exhausted", PublicError::RateLimit)] + #[case::out_of_capacity( + "429 Unable to submit request because the service is temporarily out of capacity.", + PublicError::RateLimit )] - #[case::wrapped_429( - 503, - r#"{"error": {"code": "429"}}"#, - with_debug(failure( - kept(StatusClass::RateLimit, r#"{"error": {"code": "429"}}"#), - r#"litellm.RateLimitError: vertex_aiException - {"error": {"code": "429"}}"#, - "vertex_ai", - )) - )] - #[case::overloaded( - 400, - "The model is overloaded.", - with_debug(failure( - kept(StatusClass::InternalServer, "The model is overloaded."), - "litellm.InternalServerError: vertex_aiException - The model is overloaded.", - "vertex_ai", - )) - )] - #[case::internal_server_text( - 400, - "500 Internal Server Error", - with_debug(failure( - kept(StatusClass::InternalServer, "500 Internal Server Error"), - "litellm.InternalServerError: vertex_aiException - 500 Internal Server Error", - "vertex_ai", - )) - )] - fn each_text_rule_maps_by_the_body( - #[case] status_code: u16, - #[case] body: &str, - #[case] expected: PublicFailure, - ) { - assert_eq!(mapped(&http(status_code, body)), Some(expected)); + #[case::internal_server_text("500 Internal Server Error", PublicError::InternalServer)] + #[case::overloaded("The model is overloaded.", PublicError::InternalServer)] + fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) { + assert_eq!(classified(Some(400), text), Some(expected)); } #[rstest::rstest] - #[case::bad_request( - 400, - with_debug(failure( - kept(StatusClass::BadRequest, "rejected"), - "Vertex_aiException BadRequestError - rejected", - "vertex_ai" - )) - )] - #[case::authentication( - 401, - failure( - kept(StatusClass::Authentication, "rejected"), - "Vertex_aiException - rejected", - "vertex_ai" - ) - )] - #[case::permission_denied( - 403, - failure( - kept(StatusClass::PermissionDenied, "rejected"), - "Vertex_aiException - rejected", - "vertex_ai" - ) - )] - #[case::not_found( - 404, - failure( - kept(StatusClass::NotFound, "rejected"), - "Vertex_aiException - rejected", - "vertex_ai" - ) - )] - #[case::request_timeout(408, failure(PublicKind::Timeout { status: None }, "Vertex_aiException - rejected", "vertex_ai"))] - #[case::rate_limited( - 429, - with_debug(failure( - kept(StatusClass::RateLimit, "rejected"), - "litellm.RateLimitError: Vertex_aiException - rejected", - "vertex_ai" - )) - )] - #[case::internal_server( - 500, - with_debug(failure( - kept(StatusClass::InternalServer, "rejected"), - "Vertex_aiException InternalServerError - rejected", - "vertex_ai" - )) - )] - #[case::bad_gateway( - 502, - failure( - PublicKind::ApiConnection, - "Vertex_aiException - rejected", - "vertex_ai" - ) - )] - #[case::service_unavailable( - 503, - failure( - kept(StatusClass::ServiceUnavailable, "rejected"), - "Vertex_aiException - rejected", - "vertex_ai" - ) - )] - fn each_status_rule_maps_by_the_status( - #[case] status_code: u16, - #[case] expected: PublicFailure, + #[case::server_error_wrapping_a_429(Some(503), Some(PublicError::RateLimit))] + #[case::lowest_server_error(Some(500), Some(PublicError::RateLimit))] + #[case::highest_server_error(Some(599), Some(PublicError::RateLimit))] + #[case::client_error(Some(400), None)] + #[case::no_status(None, None)] + fn a_wrapped_429_is_a_rate_limit_only_behind_a_server_error( + #[case] status: Option, + #[case] expected: Option, ) { - assert_eq!(mapped(&http(status_code, "rejected")), Some(expected)); - } - - #[rstest::rstest] - #[case::unmapped_status(409)] - #[case::gateway_timeout(504)] - fn statuses_without_a_rule_fall_through(#[case] status_code: u16) { - assert_eq!(mapped(&http(status_code, "rejected")), None); - } - - #[rstest::rstest] - #[case::stub_without_an_upstream_response( - OriginalException::Response { message: "got a 403".into() }, - status(StatusClass::BadRequest, Some(ResponseArg::Stub(HttpStub { status: 403, method: "POST", url: VERTEX_URL_WITH_SPACE, content: None }))) - )] - #[case::stub_for_a_synthesized_status( - OriginalException::Connection { message: "got a 403".into() }, - status(StatusClass::BadRequest, Some(ResponseArg::Stub(HttpStub { status: 403, method: "POST", url: VERTEX_URL_WITH_SPACE, content: None }))) - )] - fn the_rule_response_stays_when_there_is_no_real_upstream_response( - #[case] original: OriginalException, - #[case] kind: PublicKind, - ) { - assert_eq!(mapped(&original).map(|failure| failure.kind), Some(kind)); - } - - #[test] - fn a_synthesized_500_keeps_the_internal_server_stub() { - let original = OriginalException::Connection { - message: "refused".into(), - }; assert_eq!( - mapped(&original), - Some(with_debug(failure( - status( - StatusClass::InternalServer, - Some(ResponseArg::Stub(HttpStub { - status: 500, - method: "completion", - url: "https://github.com/BerriAI/litellm", - content: Some("refused".into()), - })) - ), - "Vertex_aiException InternalServerError - refused", - "vertex_ai" - ))) + classified(status, r#"{"error": {"code": "429"}}"#), + expected ); } #[rstest::rstest] #[case::project_before_payload_size( "Unable to find your project 400 Request payload size exceeds", - StatusClass::BadRequest, - "litellm.BadRequestError: vertex_aiException - Unable to find your project 400 Request payload size exceeds", - true + PublicError::BadRequest )] - #[case::payload_size_before_context_window( - "400 Request payload size exceeds; This model's maximum context length is 10", - StatusClass::ContextWindowExceeded, - "Vertex_aiException - 400 Request payload size exceeds; This model's maximum context length is 10", - false + #[case::context_window_before_unknown_error( + "This model's maximum context length is 10 None Unknown Error.", + PublicError::ContextWindowExceeded )] - #[case::api_key_before_forbidden( - "API key not valid. 403", - StatusClass::Authentication, - "Vertex_aiException - API key not valid. 403", - true + #[case::unknown_error_before_api_key( + "Content has no parts. API key not valid.", + PublicError::InternalServer )] - #[case::forbidden_before_blocked( - "403 The response was blocked.", - StatusClass::BadRequest, - "Vertex_aiException BadRequestError - 403 The response was blocked.", - true + #[case::api_key_before_blocked( + "API key not valid. The response was blocked.", + PublicError::Authentication )] #[case::blocked_before_quota( "The response was blocked. Resource exhausted", - StatusClass::ContentPolicyViolation, - "Vertex_aiException ContentPolicyViolationError - The response was blocked. Resource exhausted", - true + PublicError::ContentPolicyViolation )] #[case::quota_before_overloaded( "Resource exhausted The model is overloaded.", - StatusClass::RateLimit, - "litellm.RateLimitError: vertex_aiException - Resource exhausted The model is overloaded.", - true + PublicError::RateLimit )] - fn the_earlier_rule_wins_when_two_apply( - #[case] body: &str, - #[case] class: StatusClass, - #[case] message: &str, - #[case] debug: bool, - ) { - let expected = failure(kept(class, body), message, "vertex_ai"); - assert_eq!( - mapped(&http(401, body)), - Some(if debug { - with_debug(expected) - } else { - expected - }) - ); + fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) { + assert_eq!(classified(Some(400), text), Some(expected)); + } + + #[rstest::rstest] + #[case::a_403_in_the_text("got a 403 from 4031 tokens")] + #[case::python_client_crash("IndexError: list index out of range")] + #[case::unmarked("rejected")] + fn text_without_a_marker_is_left_to_the_status_table(#[case] text: &str) { + assert_eq!(classified(Some(400), text), None); } } diff --git a/litellm-rust/crates/core-utils/src/lib.rs b/litellm-rust/crates/core-utils/src/lib.rs index 0c26aa50cc3..fcb232d8980 100644 --- a/litellm-rust/crates/core-utils/src/lib.rs +++ b/litellm-rust/crates/core-utils/src/lib.rs @@ -4,7 +4,6 @@ pub mod exception_mapping_utils; pub mod get_llm_provider_logic; pub mod params; pub mod prompt_templates; -pub mod python_repr; pub mod secret_redaction; pub mod serde_compat; pub mod url_utils; diff --git a/litellm-rust/crates/core-utils/src/python_repr.rs b/litellm-rust/crates/core-utils/src/python_repr.rs deleted file mode 100644 index 7ccfb826377..00000000000 --- a/litellm-rust/crates/core-utils/src/python_repr.rs +++ /dev/null @@ -1,93 +0,0 @@ -/// `repr()` of a Python `str`: single quotes unless the text holds a single quote and no -/// double quote, with backslashes, the chosen quote and control characters escaped. -pub fn python_str_repr(value: &str) -> String { - let quote = if value.contains('\'') && !value.contains('"') { - '"' - } else { - '\'' - }; - let escaped: String = value - .chars() - .map(|character| match character { - '\\' => "\\\\".to_string(), - '\t' => "\\t".to_string(), - '\n' => "\\n".to_string(), - '\r' => "\\r".to_string(), - character if character == quote => format!("\\{character}"), - character - if (character as u32) < 0x20 || (0x7f..0xa0).contains(&(character as u32)) => - { - format!("\\x{:02x}", character as u32) - } - character => character.to_string(), - }) - .collect(); - format!("{quote}{escaped}{quote}") -} - -/// `repr()` of the Python value a JSON value decodes to. -pub fn python_value_repr(value: &serde_json::Value) -> String { - use serde_json::Value; - match value { - Value::Null => "None".to_string(), - Value::Bool(true) => "True".to_string(), - Value::Bool(false) => "False".to_string(), - Value::Number(number) => number.to_string(), - Value::String(text) => python_str_repr(text), - Value::Array(items) => format!( - "[{}]", - items - .iter() - .map(python_value_repr) - .collect::>() - .join(", ") - ), - Value::Object(fields) => format!( - "{{{}}}", - fields - .iter() - .map(|(key, value)| format!( - "{}: {}", - python_str_repr(key), - python_value_repr(value) - )) - .collect::>() - .join(", ") - ), - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::{python_str_repr, python_value_repr}; - - #[rstest::rstest] - #[case::null(json!(null), "None")] - #[case::true_(json!(true), "True")] - #[case::false_(json!(false), "False")] - #[case::integer(json!(5), "5")] - #[case::float(json!(1.5), "1.5")] - #[case::string(json!("it's"), "\"it's\"")] - #[case::list(json!(["a", 1]), "['a', 1]")] - #[case::dict(json!({"format": "native"}), "{'format': 'native'}")] - #[case::empty_list(json!([]), "[]")] - fn value_repr_matches_python(#[case] value: serde_json::Value, #[case] expected: &str) { - assert_eq!(python_value_repr(&value), expected); - } - - #[rstest::rstest] - #[case::plain("native", "'native'")] - #[case::single_quote("it's", "\"it's\"")] - #[case::both_quotes("it's \"x\"", "'it\\'s \"x\"'")] - #[case::double_quote("say \"x\"", "'say \"x\"'")] - #[case::backslash("a\\b", "'a\\\\b'")] - #[case::whitespace("a\tb\nc\rd", "'a\\tb\\nc\\rd'")] - #[case::control("a\u{1}b\u{7f}c\u{85}", "'a\\x01b\\x7fc\\x85'")] - #[case::unicode("café", "'café'")] - #[case::empty("", "''")] - fn matches_python_repr(#[case] value: &str, #[case] expected: &str) { - assert_eq!(python_str_repr(value), expected); - } -} diff --git a/litellm-rust/crates/core-utils/src/secret_redaction.rs b/litellm-rust/crates/core-utils/src/secret_redaction.rs index 32922d7430c..e3caee4799a 100644 --- a/litellm-rust/crates/core-utils/src/secret_redaction.rs +++ b/litellm-rust/crates/core-utils/src/secret_redaction.rs @@ -1,5 +1,3 @@ -use std::sync::LazyLock; - use fancy_regex::Regex; pub const REDACTED: &str = "REDACTED"; @@ -51,21 +49,32 @@ fn secret_patterns(minimum_custom_key_length: usize) -> String { .join("|") } -static SECRET_RE: LazyLock = LazyLock::new(|| { - Regex::new(&format!( - "(?i){}", - secret_patterns(minimum_custom_key_length()) - )) - .expect("secret redaction patterns compile") -}); - -pub fn redact_string(value: &str) -> String { - SECRET_RE.replace_all(value, REDACTED).into_owned() +/// Python's `_ENABLE_SECRET_REDACTION` pattern set, compiled once per configuration. +#[derive(Clone, Debug)] +pub struct SecretRedactor { + pattern: Regex, } -pub fn secret_redaction_enabled() -> bool { - !std::env::var("LITELLM_DISABLE_REDACT_SECRETS") - .is_ok_and(|value| value.eq_ignore_ascii_case("true")) +impl SecretRedactor { + pub fn new(minimum_custom_key_length: usize) -> Self { + let pattern = Regex::new(&format!( + "(?i){}", + secret_patterns(minimum_custom_key_length) + )) + .expect("secret redaction patterns compile"); + Self { pattern } + } + + /// `None` when `LITELLM_DISABLE_REDACT_SECRETS` turns redaction off. + pub fn from_env() -> Option { + let disabled = std::env::var("LITELLM_DISABLE_REDACT_SECRETS") + .is_ok_and(|value| value.eq_ignore_ascii_case("true")); + (!disabled).then(|| Self::new(minimum_custom_key_length())) + } + + pub fn redact(&self, value: &str) -> String { + self.pattern.replace_all(value, REDACTED).into_owned() + } } #[cfg(test)] @@ -85,13 +94,16 @@ mod tests { #[case::password_needs_word_boundary("db_password=hunter2", "REDACTED")] #[case::plain_text_is_kept(r#"{"message": "rejected"}"#, r#"{"message": "rejected"}"#)] fn redacts_the_same_spans_as_the_python_patterns(#[case] input: &str, #[case] expected: &str) { - assert_eq!(redact_string(input), expected); + assert_eq!( + SecretRedactor::new(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH).redact(input), + expected + ); } #[test] fn sk_threshold_follows_the_minimum_custom_key_length() { - let patterns = Regex::new(&format!("(?i){}", secret_patterns(8))).unwrap(); - assert_eq!(patterns.replace_all("sk-abcde", REDACTED), REDACTED); - assert_eq!(patterns.replace_all("sk-abcd", REDACTED), "sk-abcd"); + let redactor = SecretRedactor::new(8); + assert_eq!(redactor.redact("sk-abcde"), REDACTED); + assert_eq!(redactor.redact("sk-abcd"), "sk-abcd"); } } diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/api.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/api.json deleted file mode 100644 index ae629114589..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/api.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "kind": { - "type": "api", - "status": 409, - "request_url": "https://docs.litellm.ai/docs" - }, - "message": "MistralException - api", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": null, - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/api_connection.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/api_connection.json deleted file mode 100644 index ab400ed9e01..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/api_connection.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "kind": { - "type": "api_connection" - }, - "message": "MistralException - api_connection", - "model": "ocr-model", - "llm_provider": null, - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": null, - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_authentication.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_authentication.json deleted file mode 100644 index 057392575a1..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_authentication.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "authentication", - "response": null - }, - "message": "MistralException - status_authentication", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": null, - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_gateway.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_gateway.json deleted file mode 100644 index abb1425f686..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_gateway.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "bad_gateway", - "response": { - "type": "upstream", - "status": 502, - "body": "{\"message\": \"rejected\"}", - "headers": [ - [ - "retry-after", - "7" - ] - ] - } - }, - "message": "MistralException - status_bad_gateway", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": [ - [ - "retry-after", - "7" - ] - ], - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_request.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_request.json deleted file mode 100644 index 171a994cd35..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_bad_request.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "bad_request", - "response": { - "type": "upstream", - "status": 400, - "body": "{\"message\": \"rejected\"}", - "headers": [ - [ - "retry-after", - "7" - ] - ] - } - }, - "message": "MistralException - status_bad_request", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": [ - [ - "retry-after", - "7" - ] - ], - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_content_policy_violation.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_content_policy_violation.json deleted file mode 100644 index ae9c1e145de..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_content_policy_violation.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "content_policy_violation", - "response": { - "type": "upstream", - "status": 400, - "body": "{\"message\": \"rejected\"}", - "headers": [ - [ - "retry-after", - "7" - ] - ] - } - }, - "message": "MistralException - status_content_policy_violation", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": [ - [ - "retry-after", - "7" - ] - ], - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_context_window_exceeded.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_context_window_exceeded.json deleted file mode 100644 index 61e1a56a622..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_context_window_exceeded.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "context_window_exceeded", - "response": { - "type": "upstream", - "status": 400, - "body": "{\"message\": \"rejected\"}", - "headers": [ - [ - "retry-after", - "7" - ] - ] - } - }, - "message": "MistralException - status_context_window_exceeded", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": [ - [ - "retry-after", - "7" - ] - ], - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_internal_server.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_internal_server.json deleted file mode 100644 index b3c5c51a785..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_internal_server.json +++ /dev/null @@ -1,19 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "internal_server", - "response": { - "type": "stub", - "status": 500, - "method": "completion", - "url": "https://github.com/BerriAI/litellm", - "content": "upstream text" - } - }, - "message": "MistralException - status_internal_server", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": null, - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_not_found.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_not_found.json deleted file mode 100644 index 26ffd872961..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_not_found.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "not_found", - "response": { - "type": "upstream", - "status": 404, - "body": "{\"message\": \"rejected\"}", - "headers": [ - [ - "retry-after", - "7" - ] - ] - } - }, - "message": "MistralException - status_not_found", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": [ - [ - "retry-after", - "7" - ] - ], - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_permission_denied.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_permission_denied.json deleted file mode 100644 index 42772f98830..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_permission_denied.json +++ /dev/null @@ -1,19 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "permission_denied", - "response": { - "type": "stub", - "status": 403, - "method": "POST", - "url": " https://cloud.google.com/vertex-ai/", - "content": null - } - }, - "message": "MistralException - status_permission_denied", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": null, - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_rate_limit.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_rate_limit.json deleted file mode 100644 index c9b88822ebd..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_rate_limit.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "rate_limit", - "response": { - "type": "upstream", - "status": 429, - "body": "{\"message\": \"rejected\"}", - "headers": [ - [ - "retry-after", - "7" - ] - ] - } - }, - "message": "MistralException - status_rate_limit", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": [ - [ - "retry-after", - "7" - ] - ], - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_service_unavailable.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_service_unavailable.json deleted file mode 100644 index 5af65202ef1..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_service_unavailable.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "service_unavailable", - "response": { - "type": "upstream", - "status": 503, - "body": "{\"message\": \"rejected\"}", - "headers": [ - [ - "retry-after", - "7" - ] - ] - } - }, - "message": "MistralException - status_service_unavailable", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": [ - [ - "retry-after", - "7" - ] - ], - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_unsupported_params.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/status_unsupported_params.json deleted file mode 100644 index a1b318ce03c..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/status_unsupported_params.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "kind": { - "type": "status", - "status_class": "unsupported_params", - "response": { - "type": "upstream", - "status": 400, - "body": "{\"message\": \"rejected\"}", - "headers": [ - [ - "retry-after", - "7" - ] - ] - } - }, - "message": "MistralException - status_unsupported_params", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": [ - [ - "retry-after", - "7" - ] - ], - "print_banner": false -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_with_status.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_with_status.json deleted file mode 100644 index 8b215bdb7e6..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_with_status.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "kind": { - "type": "timeout", - "status": 504 - }, - "message": "MistralException - timeout_with_status", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": "\nModel: ocr-model", - "litellm_response_headers": null, - "print_banner": true -} diff --git a/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_without_status.json b/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_without_status.json deleted file mode 100644 index 4a79169776e..00000000000 --- a/tests/test_litellm/rust_bridge/fixtures/public_failures/timeout_without_status.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "kind": { - "type": "timeout", - "status": null - }, - "message": "MistralException - timeout_without_status", - "model": "ocr-model", - "llm_provider": "mistral", - "litellm_debug_info": null, - "litellm_response_headers": null, - "print_banner": false -} From 9d134413c9b6ad28b693d844a99ac1f167599f8d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 14:00:47 -0700 Subject: [PATCH 068/102] test(rust): rename the standalone 429 test to match the rule --- .../crates/core-utils/src/exception_mapping_utils/openai.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs index b4f4496dbf0..d45078c415e 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs @@ -183,7 +183,7 @@ mod tests { } #[test] - fn a_standalone_429_counts_only_with_a_429_status() { + fn a_standalone_429_counts_with_a_429_status() { assert_eq!( classified(Some(429), "got 429 back"), Some(PublicError::RateLimit) From 3d3a46fc1fcd66549958807360ad4b182117a673 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:08:11 -0700 Subject: [PATCH 069/102] refactor: build the Converse beta list once and keep the OpenAPI snapshot as CI generates it --- .../adapters/transformation.py | 18 +++++++++------- .../bedrock/chat/converse_transformation.py | 21 ++++++++++--------- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index d1047c8d86c..b213b5e30ba 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -198,6 +198,15 @@ def target_supports_mid_conversation_system(model: str | None, custom_llm_provid return supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider) +def _chat_tool_param(function_chunk: ChatCompletionToolParamFunctionChunk, tool: object) -> ChatCompletionToolParam: + eager_input_streaming: Final = eager_input_streaming_flag(tool) + if eager_input_streaming is None: + return ChatCompletionToolParam(type="function", function=function_chunk) + return ChatCompletionToolParam( + type="function", function=function_chunk, eager_input_streaming=eager_input_streaming + ) + + class AnthropicAdapter: def __init__(self) -> None: pass @@ -810,14 +819,7 @@ class LiteLLMAnthropicMessagesAdapter: for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs function_chunk.setdefault("parameters", {}).update({k: v}) - eager_input_streaming = eager_input_streaming_flag(tool) - tool_param = ( - ChatCompletionToolParam(type="function", function=function_chunk) - if eager_input_streaming is None - else ChatCompletionToolParam( - type="function", function=function_chunk, eager_input_streaming=eager_input_streaming - ) - ) + tool_param = _chat_tool_param(function_chunk, tool) self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 87e28054817..1a32fec45e3 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1518,9 +1518,6 @@ class AmazonConverseConfig(BaseConfig): """Process tools and collect anthropic_beta values.""" bedrock_tools: list[ToolBlock] = [] - # Collect anthropic_beta values from user headers - anthropic_beta_list: Final = list(get_anthropic_beta_from_headers(headers or {})) - # Separate pre-formatted Bedrock tools (e.g. systemTool from web_search_options) # from OpenAI-format tools that need transformation via _bedrock_tools_pt filtered_tools: Final = [] @@ -1540,6 +1537,17 @@ class AmazonConverseConfig(BaseConfig): continue filtered_tools.append(tool) + base_model: Final = BedrockModelInfo.get_base_model(model) + client_beta_list: Final = get_anthropic_beta_from_headers(headers or {}) + eager_beta: Final = ( + (ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER,) + if base_model.startswith("anthropic") + and AnthropicModelInfo().is_eager_input_streaming_used(filtered_tools) + and ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER not in client_beta_list + else () + ) + anthropic_beta_list: Final = [*client_beta_list, *eager_beta] + # Only separate tools if computer use tools are actually present if filtered_tools and self.is_computer_use_tool_used(filtered_tools, model): # Separate computer use tools from regular function tools @@ -1617,7 +1625,6 @@ class AmazonConverseConfig(BaseConfig): # Opus 4.5 gates ``output_config.effort`` behind a beta header; # Claude 4.6/4.7 accept it without one. - base_model: Final = BedrockModelInfo.get_base_model(model) if base_model.startswith("anthropic"): output_config: Final = additional_request_params.get("output_config") if ( @@ -1632,12 +1639,6 @@ class AmazonConverseConfig(BaseConfig): if ANTHROPIC_EFFORT_BETA_HEADER not in anthropic_beta_list: anthropic_beta_list.append(ANTHROPIC_EFFORT_BETA_HEADER) - if ( - AnthropicModelInfo().is_eager_input_streaming_used(filtered_tools) - and ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER not in anthropic_beta_list - ): - anthropic_beta_list.append(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) - # Bedrock Converse: compact_20260112 edits only (+ beta header). AmazonConverseConfig._filter_context_management_for_bedrock_converse( additional_request_params, anthropic_beta_list diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 64daa695316..7a4c966a628 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19370,7 +19370,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 93d61abfa58e52d7f4b8154ce56420522be7eb48 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 21:08:14 +0000 Subject: [PATCH 070/102] fix(deepgram): refuse /listen sessions that have no streaming price A caller could pick a model with only a pre-recorded registry row, or no row at all, and the session would be billed at the pre-recorded rate or logged at zero cost, so budgets did not apply. The route now closes the WebSocket with 1008 before dialing Deepgram unless deepgram/streaming/ (or the -multilingual row for language=multi) is an exact registry hit, and the logging handler applies the same check so a registry change under a live session records the duration with no cost instead of a substitute rate Regression tests cover the route refusal, an operator-supplied streaming row for another model being accepted, and the handler never substituting the pre-recorded rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/deepgram/common_utils.py | 39 ++++++++----- .../llm_passthrough_endpoints.py | 21 +++++-- ...gram_listen_passthrough_logging_handler.py | 19 +++--- .../deepgram/test_deepgram_common_utils.py | 58 ++++++++++++++----- ...gram_listen_passthrough_logging_handler.py | 34 +++++------ .../test_deepgram_ws_passthrough_routes.py | 57 ++++++++++++++++++ 6 files changed, 165 insertions(+), 63 deletions(-) diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index ede5b4f157d..676391dc744 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -6,8 +6,10 @@ from urllib.parse import parse_qs, urlparse import httpx +import litellm from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.utils import LlmProviders _WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"}) DEEPGRAM_LISTEN_CALLBACK_PARAMS: Final = frozenset({"callback", "callback_method"}) @@ -57,19 +59,30 @@ def _param_enabled(values: Sequence[str]) -> bool: return any(value.strip().lower() not in _DISABLED_PARAM_VALUES for value in values) -def deepgram_listen_base_pricing_models(upstream_url: str) -> tuple[str, ...]: - """Registry keys to try, in order, for the per-second base rate of a streaming session: the streaming entry for - the language mode Deepgram bills (multilingual when ``language=multi``), then the plain streaming entry, then - the pre-recorded entry for models that have no streaming price of their own.""" - model: Final = deepgram_listen_model(upstream_url) - params: Final = parse_qs(urlparse(upstream_url).query) - streaming: Final = f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{model}" - multilingual: Final = params.get("language", ("",))[-1].strip().lower() == DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE - return ( - (f"{streaming}{DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX}", streaming, model) - if multilingual - else (streaming, model) - ) +def deepgram_listen_pricing_model(upstream_url: str) -> str: + """Registry key, without the provider prefix, for the per-second base rate Deepgram bills a streaming session at: + the multilingual streaming entry when ``language=multi``, otherwise the model's own streaming entry. Pre-recorded + entries are never a substitute: Deepgram prices the two products differently.""" + streaming: Final = f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{deepgram_listen_model(upstream_url)}" + language: Final = parse_qs(urlparse(upstream_url).query).get("language", ("",))[-1] + if language.strip().lower() == DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE: + return f"{streaming}{DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX}" + return streaming + + +def deepgram_listen_registry_key(upstream_url: str) -> str: + return f"{LlmProviders.DEEPGRAM.value}/{deepgram_listen_pricing_model(upstream_url)}" + + +def deepgram_listen_is_priced(upstream_url: str) -> bool: + """Only an exact registry hit counts: the cost calculator resolves a missing ``streaming/`` row to the + pre-recorded ```` row, which is not the rate Deepgram bills a WebSocket session at.""" + registry_key: Final = deepgram_listen_registry_key(upstream_url) + try: + model_info: Final = litellm.get_model_info(model=registry_key, custom_llm_provider=LlmProviders.DEEPGRAM.value) + except Exception: + return False + return model_info["key"] == registry_key def deepgram_listen_addon_pricing_models(upstream_url: str) -> tuple[str, ...]: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ea90648e526..629784b3cee 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -38,6 +38,8 @@ from litellm.llms.azure.passthrough.transformation import foreign_azure_deployme from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.deepgram.common_utils import ( deepgram_listen_callback_params, + deepgram_listen_is_priced, + deepgram_listen_registry_key, deepgram_listen_requested_model, deepgram_listen_websocket_target, ) @@ -2897,6 +2899,9 @@ _DEEPGRAM_WS_MISSING_KEY_REASON: Final = ( "Required 'DEEPGRAM_API_KEY' in environment to make pass-through calls to Deepgram." ) _DEEPGRAM_WS_CALLBACK_REASON: Final = "Deepgram callback delivery is not supported through the proxy: remove {params}" +_DEEPGRAM_WS_UNPRICED_REASON: Final = ( + "No streaming price for '{registry_key}': add it to the model cost map to enable it" +) async def deepgram_listen_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth: @@ -2929,12 +2934,20 @@ async def deepgram_listen_websocket_route( ) return + target: Final = deepgram_listen_websocket_target( + api_base=get_secret_str("DEEPGRAM_API_BASE"), + query_string=websocket.url.query, + ) + if not deepgram_listen_is_priced(target): + await websocket.close( + code=1008, + reason=_DEEPGRAM_WS_UNPRICED_REASON.format(registry_key=deepgram_listen_registry_key(target)), + ) + return + await relay( websocket=websocket, - target=deepgram_listen_websocket_target( - api_base=get_secret_str("DEEPGRAM_API_BASE"), - query_string=websocket.url.query, - ), + target=target, custom_headers={ # mutable-ok: websocket_passthrough_request requires a plain dict of upstream headers "Authorization": f"Token {deepgram_api_key}" }, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py index c0953ae7b85..8386c154600 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/deepgram_listen_passthrough_logging_handler.py @@ -9,9 +9,11 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.deepgram.common_utils import ( deepgram_listen_addon_pricing_models, deepgram_listen_audio_seconds, - deepgram_listen_base_pricing_models, deepgram_listen_channel_count, + deepgram_listen_is_priced, deepgram_listen_model, + deepgram_listen_pricing_model, + deepgram_listen_registry_key, deepgram_listen_transcript, ) from litellm.proxy._types import PassThroughEndpointLoggingTypedDict @@ -34,19 +36,14 @@ def _registry_cost(response: TranscriptionResponse, pricing_model: str) -> float def _audio_cost(response: TranscriptionResponse, upstream_url: str) -> float | None: - base_cost: Final = next( - ( - cost - for pricing_model in deepgram_listen_base_pricing_models(upstream_url) - if (cost := _registry_cost(response, pricing_model)) is not None - ), - None, - ) - if base_cost is None: + if not deepgram_listen_is_priced(upstream_url): verbose_proxy_logger.warning( - "Deepgram listen passthrough: no pricing for model '%s'", deepgram_listen_model(upstream_url) + "Deepgram listen passthrough: no registry entry '%s'", deepgram_listen_registry_key(upstream_url) ) return None + base_cost: Final = _registry_cost(response, deepgram_listen_pricing_model(upstream_url)) + if base_cost is None: + return None addon_costs: Final = tuple( _registry_cost(response, pricing_model) for pricing_model in deepgram_listen_addon_pricing_models(upstream_url) ) diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py index 8b61d5bffa0..a1fa8f26b70 100644 --- a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py @@ -8,10 +8,12 @@ import litellm from litellm.llms.deepgram.common_utils import ( deepgram_listen_addon_pricing_models, deepgram_listen_audio_seconds, - deepgram_listen_base_pricing_models, deepgram_listen_callback_params, deepgram_listen_channel_count, + deepgram_listen_is_priced, deepgram_listen_model, + deepgram_listen_pricing_model, + deepgram_listen_registry_key, deepgram_listen_requested_model, deepgram_listen_transcript, deepgram_listen_websocket_target, @@ -220,27 +222,51 @@ def test_requested_model_is_the_model_the_upstream_target_will_carry(query_strin @pytest.mark.parametrize( ("upstream_url", "expected"), [ - pytest.param(NOVA_3_URL, ("streaming/nova-3", "nova-3"), id="monolingual"), - pytest.param(f"{NOVA_3_URL}&language=en", ("streaming/nova-3", "nova-3"), id="explicit language"), - pytest.param( - f"{NOVA_3_URL}&language=multi", - ("streaming/nova-3-multilingual", "streaming/nova-3", "nova-3"), - id="multilingual", - ), - pytest.param( - f"{NOVA_3_URL}&language=MULTI", - ("streaming/nova-3-multilingual", "streaming/nova-3", "nova-3"), - id="multilingual any case", - ), + pytest.param(NOVA_3_URL, "streaming/nova-3", id="monolingual"), + pytest.param(f"{NOVA_3_URL}&language=en", "streaming/nova-3", id="explicit language"), + pytest.param(f"{NOVA_3_URL}&language=multi", "streaming/nova-3-multilingual", id="multilingual"), + pytest.param(f"{NOVA_3_URL}&language=MULTI", "streaming/nova-3-multilingual", id="multilingual any case"), pytest.param( "wss://api.deepgram.com/v1/listen?model=nova-2&language=multi", - ("streaming/nova-2-multilingual", "streaming/nova-2", "nova-2"), + "streaming/nova-2-multilingual", id="other model", ), + pytest.param("wss://api.deepgram.com/v1/listen?encoding=linear16", "streaming/nova-3", id="default model"), ], ) -def test_deepgram_listen_base_pricing_models(upstream_url: str, expected: tuple[str, ...]): - assert deepgram_listen_base_pricing_models(upstream_url) == expected +def test_deepgram_listen_pricing_model_is_the_streaming_entry_never_the_prerecorded_one( + upstream_url: str, expected: str +): + assert deepgram_listen_pricing_model(upstream_url) == expected + assert deepgram_listen_registry_key(upstream_url) == f"deepgram/{expected}" + + +NOVA_2_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-2" + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("upstream_url", "extra_rows", "expected"), + [ + pytest.param(NOVA_3_URL, (), True, id="streaming entry present"), + pytest.param(f"{NOVA_3_URL}&language=multi", (), True, id="multilingual entry present"), + pytest.param(NOVA_2_URL, (), False, id="only the pre-recorded entry"), + pytest.param(f"{NOVA_2_URL}&language=multi", ("deepgram/streaming/nova-2",), False, id="needs multilingual"), + pytest.param("wss://api.deepgram.com/v1/listen?model=nova-99-unmapped", (), False, id="nothing priced"), + pytest.param(NOVA_2_URL, ("deepgram/streaming/nova-2",), True, id="operator-supplied streaming entry"), + pytest.param(NOVA_2_URL, ("streaming/nova-2",), False, id="a row under another key is not the entry"), + ], +) +def test_deepgram_listen_is_priced( + monkeypatch: pytest.MonkeyPatch, upstream_url: str, extra_rows: tuple[str, ...], expected: bool +): + """The bundled map prices only nova-3 for streaming; nova-2 has a pre-recorded row, which must never count.""" + monkeypatch.delitem(litellm.model_cost, "deepgram/streaming/nova-2", raising=False) + assert "deepgram/nova-2" in litellm.model_cost + for row in extra_rows: + monkeypatch.setitem(litellm.model_cost, row, dict(litellm.model_cost["deepgram/streaming/nova-3"])) + + assert deepgram_listen_is_priced(upstream_url) is expected @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py index 2ae56709b11..ac742c0ab46 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_listen_passthrough_logging_handler.py @@ -158,17 +158,25 @@ def test_handler_add_ons_scale_with_channels_like_the_base_rate(): assert stereo_redacted - stereo_plain == pytest.approx(_registry_cost("streaming/redact", 120.0)) -def test_handler_falls_back_to_the_prerecorded_rate_for_a_model_without_a_streaming_entry(): - assert "deepgram/streaming/nova-2" not in litellm.model_cost +@pytest.mark.parametrize( + "upstream_url", + [ + pytest.param("wss://api.deepgram.com/v1/listen?model=nova-2", id="only a pre-recorded entry"), + pytest.param("wss://api.deepgram.com/v1/listen?model=nova-99-not-in-registry", id="no entry at all"), + ], +) +def test_handler_never_substitutes_another_rate_for_a_missing_streaming_entry(monkeypatch, upstream_url): + """The route refuses these sessions up front; should the registry change under a live one, the spend row + keeps the duration and carries no cost, rather than the pre-recorded rate or any other stand-in.""" + monkeypatch.delitem(litellm.model_cost, "deepgram/streaming/nova-2", raising=False) + assert "deepgram/nova-2" in litellm.model_cost handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( - websocket_messages=(_metadata(60.0),), - logging_obj=_logging_obj(), - upstream_url="wss://api.deepgram.com/v1/listen?model=nova-2", + websocket_messages=(_metadata(60.0),), logging_obj=_logging_obj(), upstream_url=upstream_url ) - assert handler_result["kwargs"]["model"] == "nova-2" - assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("nova-2", 60.0)) + assert handler_result["kwargs"]["response_cost"] is None + assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 60.0 def test_handler_falls_back_to_results_frames_when_the_stream_ends_without_metadata(): @@ -222,18 +230,6 @@ def test_handler_bills_the_declared_channels_when_the_stream_dies_before_any_fra assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 30.0)) -def test_handler_keeps_the_spend_row_but_no_cost_for_an_unpriced_model(): - handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( - websocket_messages=(_metadata(12.5),), - logging_obj=_logging_obj(), - upstream_url="wss://api.deepgram.com/v1/listen?model=nova-99-not-in-registry", - ) - - assert handler_result["kwargs"]["model"] == "nova-99-not-in-registry" - assert handler_result["kwargs"]["response_cost"] is None - assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 12.5 - - class _CapturingLogger(CustomLogger): def __init__(self) -> None: super().__init__() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 85baaacf1e9..4eb183b14ce 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -30,6 +30,14 @@ GET_CREDENTIALS: Final = ( ) USER_API_KEY_AUTH: Final = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" LISTEN_PATHS: Final = ("/deepgram/v1/listen", "/deepgram/listen") +NOVA_2_STREAMING_KEY: Final = "deepgram/streaming/nova-2" + +pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map") + + +def _price_nova_2_streaming(monkeypatch: pytest.MonkeyPatch) -> None: + """An operator-supplied streaming row: the bundled map prices only nova-3 for streaming.""" + monkeypatch.setitem(litellm.model_cost, NOVA_2_STREAMING_KEY, dict(litellm.model_cost["deepgram/streaming/nova-3"])) class _FakeWebSocket: @@ -138,6 +146,7 @@ async def test_deepgram_listen_forwards_query_and_injects_only_provider_auth(pat @pytest.mark.asyncio async def test_deepgram_listen_keeps_caller_chosen_model(monkeypatch): monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + _price_nova_2_streaming(monkeypatch) websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-2&language=en") with patch(GET_CREDENTIALS, return_value="dg-provider-key"): @@ -236,6 +245,53 @@ async def test_deepgram_listen_rejects_callback_delivery_that_would_go_unbilled( assert "dg-provider-key" not in websocket.closed[1] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("query", "missing_key"), + [ + pytest.param("model=nova-2", "deepgram/streaming/nova-2", id="model with only a pre-recorded price"), + pytest.param("model=nova-99", "deepgram/streaming/nova-99", id="model unknown to the registry"), + pytest.param( + "model=nova-3&language=multi", + "deepgram/streaming/nova-3-multilingual", + id="multilingual session without its own price", + ), + ], +) +async def test_deepgram_listen_refuses_sessions_it_cannot_price(query, missing_key, monkeypatch): + """A session with no streaming price would be logged at zero (or at the pre-recorded rate), letting a caller run + up unmetered spend, so the proxy closes it before Deepgram is contacted and names the registry row to add.""" + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.delitem(litellm.model_cost, missing_key, raising=False) + assert "deepgram/nova-2" in litellm.model_cost + websocket = _FakeWebSocket("/deepgram/v1/listen", query) + + with patch(GET_CREDENTIALS, return_value="dg-provider-key"): + relay = await _serve(websocket) + + assert relay.calls == [] + assert websocket.closed is not None + assert websocket.closed[0] == 1008 + assert missing_key in websocket.closed[1] + assert "dg-provider-key" not in websocket.closed[1] + + +@pytest.mark.asyncio +async def test_deepgram_listen_relays_once_the_operator_prices_the_model(monkeypatch): + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-2") + with patch(GET_CREDENTIALS, return_value="dg-provider-key"): + assert (await _serve(websocket)).calls == [] + + _price_nova_2_streaming(monkeypatch) + priced_websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-2") + with patch(GET_CREDENTIALS, return_value="dg-provider-key"): + relay = await _serve(priced_websocket) + + assert [call.target for call in relay.calls] == ["wss://api.deepgram.com/v1/listen?model=nova-2"] + assert priced_websocket.closed is None + + def _app_with_relay(relay: _FakeRelay) -> FastAPI: app = FastAPI() app.include_router(router) @@ -332,6 +388,7 @@ def test_deepgram_listen_authorizes_the_model_it_will_actually_send_upstream(que in its default: the real key auth path must see the same model the upstream target will carry.""" monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) monkeypatch.setattr(litellm, "max_budget", 0.0) + _price_nova_2_streaming(monkeypatch) cache = asyncio.run(_cache_restricted_key("sk-only-nova-2", ["nova-2"])) relay = _FakeRelay() client = TestClient(_app_with_relay(relay)) From f1b9642c4108e5bb3c9fbae858d1f39099f815cb Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 21:15:01 +0000 Subject: [PATCH 071/102] fix(proxy): classify Azure Speech short audio behind a prefixed api base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../azure_speech_passthrough_logging_handler.py | 4 +++- .../test_azure_speech_passthrough_logging_handler.py | 8 ++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/azure_speech_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/azure_speech_passthrough_logging_handler.py index 8cd1b137de8..33d1815b3c4 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/azure_speech_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/azure_speech_passthrough_logging_handler.py @@ -8,6 +8,7 @@ import httpx from litellm._logging import verbose_proxy_logger from litellm.constants import ( AZURE_SPEECH_BATCH_MODEL, + AZURE_SPEECH_BATCH_PATH_PREFIX, AZURE_SPEECH_CUSTOM_LLM_PROVIDER, AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL, AZURE_SPEECH_FAST_TRANSCRIPTION_PATH, @@ -30,7 +31,8 @@ from litellm.types.utils import StandardPassThroughResponseObject class AzureSpeechPassthroughLoggingHandler: @staticmethod def _is_short_audio_route(url_route: str) -> bool: - return urlparse(url_route).path.startswith(AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX) + path: Final = urlparse(url_route).path + return path.rfind(AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX) > path.rfind(AZURE_SPEECH_BATCH_PATH_PREFIX) @staticmethod def _is_fast_transcription_route(url_route: str) -> bool: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_azure_speech_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_azure_speech_passthrough_logging_handler.py index 670f8e65823..a0d27e618f9 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_azure_speech_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_azure_speech_passthrough_logging_handler.py @@ -21,6 +21,11 @@ SHORT_AUDIO_URL = ( ) BATCH_URL = "https://eastus.api.cognitive.microsoft.com/speechtotext/v3.2/transcriptions" FAST_URL = "https://eastus.api.cognitive.microsoft.com/speechtotext/transcriptions:transcribe?api-version=2024-11-15" +PREFIXED_SHORT_AUDIO_URL = ( + "https://apim.example.com/speech-proxy/speech/recognition/conversation/cognitiveservices/v1?language=en-US" +) +PREFIXED_BATCH_URL = "https://apim.example.com/speech/speechtotext/v3.2/transcriptions" +PREFIXED_FAST_URL = "https://apim.example.com/speech/speechtotext/transcriptions:transcribe?api-version=2024-11-15" FAST_BODY = {"durationMilliseconds": 5061, "combinedPhrases": [{"text": "Hello world."}]} FAST_AUDIO_SECONDS = 5.061 TRANSCRIPT_BODY = { @@ -87,6 +92,9 @@ class TestAzureSpeechPassthroughHandler: (FAST_URL, "azure_speech/fast-transcription", FAST_AUDIO_SECONDS * PRICE_PER_SECOND), (BATCH_URL, "azure_speech/batch-transcription", 0.0), (f"{BATCH_URL}/8a5d3f2c-0b1e-4c7d-9e6f-1234567890ab/files", "azure_speech/batch-transcription", 0.0), + (PREFIXED_SHORT_AUDIO_URL, "azure_speech/short-audio", TRANSCRIPT_AUDIO_SECONDS * PRICE_PER_SECOND), + (PREFIXED_FAST_URL, "azure_speech/fast-transcription", FAST_AUDIO_SECONDS * PRICE_PER_SECOND), + (PREFIXED_BATCH_URL, "azure_speech/batch-transcription", 0.0), ], ) def test_records_model_provider_and_cost(self, url_route: str, expected_model: str, expected_cost: float): From 3449ae9d0d1d339146fa5ffb7318a62553860a46 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 21:15:32 +0000 Subject: [PATCH 072/102] fix(proxy): advance the daily global spend marker in one conditional upsert so overlapping runs cannot rewind it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../daily_global_spend_rollup.py | 51 +++++++--------- .../test_daily_global_spend_rollup.py | 61 ++++++++++++++++--- 2 files changed, 74 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index d50feb6f16b..376b113ed02 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -23,7 +23,6 @@ from litellm.constants import ( DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ) -from litellm.repositories.config_repository import ConfigRepository if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache @@ -82,6 +81,17 @@ _PENDING_DAYS_SQL: Final = ( 'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') ' 'ORDER BY "date"' ) +# Runs can overlap (Redis unreachable, lock expired on a long backfill), so the database keeps the +# later of the stored and the incoming day and scan time in one statement; GREATEST skips NULL. +_ADVANCE_MARKER_SQL: Final = ( + 'INSERT INTO "LiteLLM_Config" ("param_name", "param_value") ' + "VALUES ($1, jsonb_build_object('reconciled_through', $2::text, 'scanned_at', $3::text)) " + 'ON CONFLICT ("param_name") DO UPDATE SET "param_value" = jsonb_build_object(' + "'reconciled_through', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'reconciled_through', " + "EXCLUDED.\"param_value\" ->> 'reconciled_through'), " + "'scanned_at', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'scanned_at', " + "EXCLUDED.\"param_value\" ->> 'scanned_at'))" +) class ReconciledThrough(BaseModel): @@ -152,33 +162,20 @@ async def reconciled_through(prisma_client: "PrismaClient") -> str | None: return None if marker is None else marker.reconciled_through -async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThrough) -> None: +async def _advance_marker(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: + """Move the stored marker to the last of ``days`` and to ``scanned_at`` where those are later + than what is stored, so a slower overlapping run can only add to a faster run's marker.""" from litellm.proxy.utils import invalidate_config_param - await ConfigRepository(prisma_client).set_param( - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, marker.model_dump_json() + await prisma_client.db.execute_raw( + _ADVANCE_MARKER_SQL, + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, + max(days) if days else None, + scanned_at, ) await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -async def _stored_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: - """The marker as another pod may have just written it, bypassing this pod's config cache.""" - param: Final = await ConfigRepository(prisma_client).get_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - return None if param is None else _marker_from_param_value(param.param_value) - - -async def _record_advanced(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: - """Advance the stored marker by ``days``. Two runs can overlap (Redis unreachable, lock expired - on a long backfill), so the base is what is stored now, not the snapshot this run scanned from: - a slower run may then only add to the faster run's marker, never rewind it. Without a new scan - time the stored one is kept.""" - stored: Final = await _stored_marker(prisma_client) - kept_scanned_at: Final = None if stored is None else stored.scanned_at - await _record_marker( - prisma_client, _advanced(stored, days, scanned_at=scanned_at if scanned_at is not None else kept_scanned_at) - ) - - async def _db_now(prisma_client: "PrismaClient") -> _NowRow: rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) return _NowRow.model_validate(rows[0]) @@ -221,16 +218,10 @@ async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> Rec marker: Final = await reconciled_through(prisma_client) return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) if scan.marker is not None or done: - await _record_advanced(prisma_client, done, scanned_at=scan.scanned_at) + await _advance_marker(prisma_client, done, scanned_at=scan.scanned_at) return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) -def _advanced(marker: ReconciledThrough | None, days: tuple[str, ...], *, scanned_at: str | None) -> ReconciledThrough: - """The marker after ``days`` were rewritten: a late old day never moves it back.""" - through: Final = max((marker.reconciled_through if marker is not None else "", *days)) - return ReconciledThrough(reconciled_through=through, scanned_at=scanned_at) - - async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: for index, day in enumerate(scan.days): if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]): @@ -242,7 +233,7 @@ async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: t day: Final = done_with_this[-1] try: await reconcile_day(prisma_client, day) - await _record_advanced(prisma_client, done_with_this, scanned_at=None) + await _advance_marker(prisma_client, done_with_this, scanned_at=None) except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) return False diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 69f06fad081..3da587435ad 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -1,5 +1,6 @@ """Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818).""" +import json import pathlib import re from datetime import date @@ -14,6 +15,7 @@ from pytest_postgresql import factories from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( + _ADVANCE_MARKER_SQL, RECONCILE_DAY_SQL, read_marker, reconciled_through, @@ -36,13 +38,21 @@ class _FakeConfigTable: def __init__(self) -> None: self.rows: dict[str, object] = {} - async def upsert(self, *, where: dict[str, str], data: dict[str, dict[str, str]]) -> _FakeConfigRow: - self.rows[where["param_name"]] = data["update"]["param_value"] - return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) + def advance(self, param_name: str, through: str | None, scanned_at: str | None) -> None: + """What ``_ADVANCE_MARKER_SQL`` does in Postgres: keep the later of stored and incoming per field.""" + stored = self.rows.get(param_name) + current: dict[str, str | None] = json.loads(stored) if isinstance(stored, str) else {} + self.rows[param_name] = json.dumps( + { + "reconciled_through": _greatest(current.get("reconciled_through"), through), + "scanned_at": _greatest(current.get("scanned_at"), scanned_at), + } + ) - async def find_unique(self, *, where: dict[str, str]) -> _FakeConfigRow | None: - stored = self.rows.get(where["param_name"]) - return None if stored is None else _FakeConfigRow(where["param_name"], stored) + +def _greatest(stored: str | None, incoming: str | None) -> str | None: + present = [value for value in (stored, incoming) if value is not None] + return max(present) if present else None class _FakeDb: @@ -67,9 +77,14 @@ class _FakeDb: {"date": d} for d, written in sorted(rows.items()) if d <= last and (d > marker or written >= scanned_at) ] - async def execute_raw(self, sql: str, *params: str) -> int: + async def execute_raw(self, sql: str, *params: str | None) -> int: + if sql == _ADVANCE_MARKER_SQL: + param_name, through, scanned_at = params + assert param_name is not None + self.litellm_config.advance(param_name, through, scanned_at) + return 1 (day,) = params - if day in self._prisma.failing_days: + if day is None or day in self._prisma.failing_days: raise RuntimeError(f"day {day} exploded") self._prisma.reconciled.append(day) landing = self._prisma.marker_landing_on_day.get(day) @@ -485,3 +500,33 @@ def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_p assert sum(int(r["total_response_time_ms"]) for r in global_rows) == 1600 # pyright: ignore[reportArgumentType] # dict_row values are untyped assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")] assert untouched == [] + + +_CONFIG_DDL: Final = 'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB)' +_MARKER_SQL: Final = 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s' + + +def test_advance_marker_sql_only_ever_moves_the_stored_marker_forward(_rollup_postgresql: psycopg.Connection): + """Against real Postgres: the statement a slower overlapping run issues after the faster run + already stored a later marker leaves that marker alone, whether it carries an older scan time or + none at all, while a run that is further along moves both fields on.""" + conn: Final = _rollup_postgresql + conn.execute(_CONFIG_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.commit() + param: Final = DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM + + def stored() -> object: + with conn.cursor(row_factory=dict_row) as cur: + row = cur.execute(_MARKER_SQL, (param,)).fetchone() + return None if row is None else row["param_value"] + + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-01", None)) + assert stored() == {"reconciled_through": "2026-09-01", "scanned_at": None} + + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-14", "2026-09-15 00:30:02.5")) + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-02", None)) + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-03", "2026-09-15 00:30:01.25")) + assert stored() == {"reconciled_through": "2026-09-14", "scanned_at": "2026-09-15 00:30:02.5"} + + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-15", "2026-09-16 00:30:00.75")) + assert stored() == {"reconciled_through": "2026-09-15", "scanned_at": "2026-09-16 00:30:00.75"} From 4a951847bbcb25a8c42eb526b0526604d1b3bb1a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:15:39 -0700 Subject: [PATCH 073/102] fix(responses): merge deployment litellm_params into native websocket response.create frames --- litellm/llms/custom_httpx/llm_http_handler.py | 3 + litellm/proxy/_lazy_openapi_snapshot.json | 2 +- litellm/responses/main.py | 25 ++- litellm/responses/streaming_iterator.py | 31 ++- .../types/responses/streaming_websocket.py | 13 ++ .../test_responses_websocket_all_providers.py | 210 ++++++++++++++++++ 6 files changed, 275 insertions(+), 9 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 98fe0014386..477d10a3cbd 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -144,6 +144,7 @@ from litellm.types.llms.openai import ( from litellm.types.realtime import RealtimeQueryParams from litellm.types.rerank import RerankResponse from litellm.types.responses.main import DeleteResponseResult +from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CallTypes, @@ -6588,6 +6589,7 @@ class BaseLLMHTTPHandler: litellm_metadata: dict[str, object] | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, + request_defaults: ResponsesWebSocketRequestDefaults | None = None, **kwargs: Any, ): """ @@ -6742,6 +6744,7 @@ class BaseLLMHTTPHandler: output_guardrail_callbacks=_ws_output_guardrail_callbacks, quota_callbacks=_ws_quota_callbacks, authorized_model=model, + request_defaults=request_defaults, ) await streaming.bidirectional_forward() diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 8a8d08c6887..2aa2cf15ac1 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19394,7 +19394,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 6dc34bb93ef..6064fa66d91 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast import httpx -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import assert_never import litellm @@ -53,6 +53,7 @@ from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * +from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import all_litellm_params from litellm.utils import ( @@ -2261,6 +2262,27 @@ def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: return metadata +_EXTRA_BODY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) + + +def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults: + reasoning_effort: Final = kwargs.get("reasoning_effort") + mapped_reasoning: Final = ( + LiteLLMResponsesTransformationHandler()._map_reasoning_effort(reasoning_effort) + if kwargs.get("reasoning") is None and isinstance(reasoning_effort, str) + else None + ) + candidate_params: Final[dict[str, object]] = { + **kwargs, + **({"reasoning": mapped_reasoning} if mapped_reasoning is not None else {}), + } + fill_missing: Final = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(candidate_params) + return ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType(dict(fill_missing)), + overrides=MappingProxyType(_EXTRA_BODY_ADAPTER.validate_python(kwargs.get("extra_body")) or {}), + ) + + @client async def _aresponses_websocket( model: str, @@ -2352,5 +2374,6 @@ async def _aresponses_websocket( user_api_key_dict=kwargs.get("user_api_key_dict"), litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, + request_defaults=_build_responses_websocket_request_defaults(kwargs), **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8d766cf1cd0..1c82df8c664 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -54,6 +54,7 @@ if TYPE_CHECKING: PresidioGuardrailCallback, ResponsesBackendWebSocket, ResponsesClientWebSocket, + ResponsesWebSocketRequestDefaults, ) from litellm.types.router import LiteLLM_Params @@ -1717,6 +1718,7 @@ class ResponsesWebSocketStreaming: output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, + request_defaults: ResponsesWebSocketRequestDefaults | None = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1732,6 +1734,7 @@ class ResponsesWebSocketStreaming: # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model + self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES @@ -1874,12 +1877,23 @@ class ResponsesWebSocketStreaming: modified = True return modified + def _with_request_defaults(self, msg_obj: dict[str, object]) -> dict[str, object]: + if self.request_defaults is None: + return msg_obj + nested: Final = msg_obj.get("response") + if _is_json_object(nested): + return {**msg_obj, "response": self.request_defaults.merged_into(nested)} + return self.request_defaults.merged_into(msg_obj) + async def _mask_response_create(self, message: str) -> str: """ - Enforce the authorized model and apply Presidio PII masking to a - ``response.create`` message before it is forwarded to the upstream - provider. + Merge deployment defaults, enforce the authorized model, and apply + Presidio PII masking to a ``response.create`` message before it is + forwarded to the upstream provider. + - Fills the deployment's ``litellm_params`` request defaults into the + frame the way the HTTP ``/v1/responses`` path does: client-set keys + win, ``extra_body`` entries override. - Overwrites any ``model`` field with the connection-authorized model to prevent deployment-substitution attacks (always applied). - Walks the ``input`` and ``instructions`` fields, calls ``check_pii`` @@ -1889,23 +1903,26 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ try: - msg_obj: Final = _load_json_object(message) + parsed: Final = _load_json_object(message) except (json.JSONDecodeError, TypeError): return message - if msg_obj.get("type") != "response.create": + if parsed.get("type") != "response.create": return message + msg_obj: Final = self._with_request_defaults(parsed) + defaults_applied: Final = msg_obj != parsed + # Always enforce the authorized model, even when PII masking is off. model_modified: Final = self._enforce_authorized_model(msg_obj) if not self.guardrail_callbacks: - return json.dumps(msg_obj) if model_modified else message + return json.dumps(msg_obj) if model_modified or defaults_applied else message if "metadata" not in self.request_data: self.request_data["metadata"] = {} - modified = model_modified + modified = model_modified or defaults_applied guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks) for cb in guardrail_cbs: presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) diff --git a/litellm/types/responses/streaming_websocket.py b/litellm/types/responses/streaming_websocket.py index 2aa71647955..f369cbcebf8 100644 --- a/litellm/types/responses/streaming_websocket.py +++ b/litellm/types/responses/streaming_websocket.py @@ -1,5 +1,7 @@ from __future__ import annotations +from collections.abc import Mapping +from dataclasses import dataclass from typing import Protocol from litellm.types.guardrails import PresidioPerRequestConfig @@ -39,3 +41,14 @@ class PresidioGuardrailCallback(Protocol): presidio_config: PresidioPerRequestConfig | None, request_data: dict[str, object], ) -> str: ... + + +@dataclass(frozen=True, slots=True) +class ResponsesWebSocketRequestDefaults: + """Deployment-level request parameters merged into every ``response.create`` frame relayed over a native websocket.""" + + fill_missing: Mapping[str, object] + overrides: Mapping[str, object] + + def merged_into(self, request: Mapping[str, object]) -> dict[str, object]: + return {**self.fill_missing, **request, **self.overrides} diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index fe3c4a0640d..918c4a33d40 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1204,6 +1204,216 @@ class TestWebSocketProjectQuotaEnforcement: quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() +def _deployment_defaults(): + from types import MappingProxyType + + from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults + + return ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType({"reasoning": {"effort": "high"}, "service_tier": "priority"}), + overrides=MappingProxyType({"provider_default": "configured"}), + ) + + +class TestNativeWebSocketDeploymentDefaults: + """The native relay merges deployment litellm_params into every response.create like HTTP does.""" + + def test_builder_maps_router_kwargs_like_the_http_path(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + { + "model": "gpt-5-pro", + "reasoning_effort": "high", + "service_tier": "priority", + "extra_body": {"provider_default": "configured"}, + "temperature": None, + "timeout": 600, + "max_retries": 2, + "caching": False, + "custom_llm_provider": "openai", + "litellm_metadata": {"user_api_key": "hashed"}, + "user_api_key_dict": MagicMock(), + "litellm_logging_obj": MagicMock(), + "websocket": MagicMock(), + } + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "high"}, "service_tier": "priority"} + assert dict(defaults.overrides) == {"provider_default": "configured"} + + def test_builder_keeps_explicit_reasoning_over_reasoning_effort(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + {"model": "gpt-5-pro", "reasoning": {"effort": "low"}, "reasoning_effort": "high"} + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "low"}} + assert dict(defaults.overrides) == {} + + @pytest.mark.asyncio + async def test_flat_frame_gets_defaults_client_keys_win_extra_body_overrides(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "service_tier": "default", + "provider_default": "client", + } + ) + ) + ) + + assert forwarded == { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "service_tier": "default", + "provider_default": "configured", + "reasoning": {"effort": "high"}, + } + + @pytest.mark.asyncio + async def test_nested_response_frame_gets_defaults_inside_response(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps({"type": "response.create", "response": {"model": "gpt-5-pro", "input": "hi"}}) + ) + ) + + assert forwarded == { + "type": "response.create", + "response": { + "model": "gpt-5-pro", + "input": "hi", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + }, + } + + @pytest.mark.asyncio + async def test_frames_that_need_nothing_pass_through_untouched(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + cancel_frame = json.dumps({"type": "response.cancel"}) + complete_frame = json.dumps( + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "hi", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + } + ) + + assert await handler._mask_response_create(cancel_frame) is cancel_frame + assert await handler._mask_response_create(complete_frame) is complete_frame + + @pytest.mark.asyncio + async def test_handler_applies_defaults_to_the_first_frame_sent_upstream(self): + import asyncio + from unittest.mock import AsyncMock, patch + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + class FakeBackend: + def __init__(self): + self.sent = [] + + async def send(self, message): + self.sent.append(message) + + async def recv(self, decode=False): + raise RuntimeError("backend closed") + + async def close(self): + pass + + backend = FakeBackend() + + class FakeConnect: + def __init__(self, url, **kwargs): + pass + + async def __aenter__(self): + return backend + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.model_in_websocket_url.return_value = True + mock_config.get_websocket_url.return_value = "wss://api.openai.com/v1/responses" + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + mock_logging.dispatch_success_handlers = AsyncMock() + + client_ws = MagicMock() + client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client closed")) + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await BaseLLMHTTPHandler().async_responses_websocket( + model="gpt-5-pro", + websocket=client_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + first_message=json.dumps({"type": "response.create", "model": "gpt-5-pro", "input": "Say hello"}), + request_defaults=_deployment_defaults(), + ) + await asyncio.sleep(0) + + assert [json.loads(frame) for frame in backend.sent] == [ + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + } + ] + + @pytest.mark.asyncio + async def test_aresponses_websocket_builds_defaults_from_deployment_kwargs(self, monkeypatch): + import importlib + from unittest.mock import AsyncMock + + responses_main = importlib.import_module("litellm.responses.main") + + stub = MagicMock() + stub.async_responses_websocket = AsyncMock() + monkeypatch.setattr(responses_main, "base_llm_http_handler", stub) + + await responses_main._aresponses_websocket.__wrapped__( + model="openai/gpt-5-pro", + websocket=MagicMock(), + api_key="sk-test", + litellm_logging_obj=MagicMock(), + reasoning_effort="high", + service_tier="priority", + extra_body={"provider_default": "configured"}, + ) + + request_defaults = stub.async_responses_websocket.call_args.kwargs["request_defaults"] + assert dict(request_defaults.fill_missing) == {"reasoning": {"effort": "high"}, "service_tier": "priority"} + assert dict(request_defaults.overrides) == {"provider_default": "configured"} + + class TestNativeWebSocketGuardrails: @pytest.mark.asyncio async def test_response_create_injects_authorized_model(self): From 1f27c442b4146e61d77c904a5c5d9e5165602267 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 18 Sep 2026 14:17:00 -0700 Subject: [PATCH 074/102] ci(duplicate-check): let Codex reach GitHub from its sandbox Every gh search in the first real runs failed with "error connecting to api.github.com", so the verdict was always null. The legacy sandbox_permissions key no longer grants network in read-only mode; the workspace-write sandbox has a network_access switch that does. Pin the CLI to the version the prompt was proven on --- .github/workflows/duplicate_issue_check.yml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/duplicate_issue_check.yml b/.github/workflows/duplicate_issue_check.yml index b12f894328e..83c6cc06de7 100644 --- a/.github/workflows/duplicate_issue_check.yml +++ b/.github/workflows/duplicate_issue_check.yml @@ -93,10 +93,11 @@ jobs: responses-api-endpoint: ${{ vars.LITELLM_API_BASE }}/v1/responses prompt-file: .github/prompts/duplicate-issue-check.md output-schema-file: .github/prompts/duplicate-issue-check.schema.json - sandbox: read-only - # read-only denies network, and the whole method is searching the tracker with gh - codex-args: '["-c", "sandbox_permissions=[\"network-full-access\"]"]' + sandbox: workspace-write + # The whole method is searching the tracker with gh, and network is only switchable in workspace-write + codex-args: '["-c", "sandbox_workspace_write.network_access=true"]' model: ${{ vars.DUPLICATE_CHECK_MODEL }} + codex-version: "0.154.0" # Issue authors have no write access and the action refuses them by default; the # prompt is fixed, the sandbox read-only, and the only token is read-only on a public repo allow-users: "*" From 1e8b8f7c33b83b35463e5efa9c11c33eb3929d54 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 18 Sep 2026 14:21:22 -0700 Subject: [PATCH 075/102] ci(duplicate-check): describe the workspace-write sandbox accurately --- .github/workflows/duplicate_issue_check.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/duplicate_issue_check.yml b/.github/workflows/duplicate_issue_check.yml index 83c6cc06de7..91eee3ca383 100644 --- a/.github/workflows/duplicate_issue_check.yml +++ b/.github/workflows/duplicate_issue_check.yml @@ -98,8 +98,8 @@ jobs: codex-args: '["-c", "sandbox_workspace_write.network_access=true"]' model: ${{ vars.DUPLICATE_CHECK_MODEL }} codex-version: "0.154.0" - # Issue authors have no write access and the action refuses them by default; the - # prompt is fixed, the sandbox read-only, and the only token is read-only on a public repo + # Issue authors have no write access and the action refuses them by default; the prompt is + # fixed, writes stay inside the throwaway checkout, and the only token is read-only on a public repo allow-users: "*" - name: Summary From 9767878425ea43788925d567ce5499cc02dc44f8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 21:27:03 +0000 Subject: [PATCH 076/102] test(ocr): replace per-provider OCR test classes with a declarative provider x auth x input matrix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/ocr_tests/base_ocr_unit_tests.py | 203 ----------- tests/ocr_tests/test_ocr_azure_ai.py | 29 -- .../test_ocr_azure_document_intelligence.py | 49 +-- tests/ocr_tests/test_ocr_matrix.py | 317 ++++++++++++++++++ tests/ocr_tests/test_ocr_mistral.py | 92 ----- tests/ocr_tests/test_ocr_vertex_ai.py | 127 +------ tests/ocr_tests/vertex_key.json | 13 - 7 files changed, 328 insertions(+), 502 deletions(-) delete mode 100644 tests/ocr_tests/base_ocr_unit_tests.py delete mode 100644 tests/ocr_tests/test_ocr_azure_ai.py create mode 100644 tests/ocr_tests/test_ocr_matrix.py delete mode 100644 tests/ocr_tests/test_ocr_mistral.py delete mode 100644 tests/ocr_tests/vertex_key.json diff --git a/tests/ocr_tests/base_ocr_unit_tests.py b/tests/ocr_tests/base_ocr_unit_tests.py deleted file mode 100644 index ae65efd952d..00000000000 --- a/tests/ocr_tests/base_ocr_unit_tests.py +++ /dev/null @@ -1,203 +0,0 @@ -""" -Base test class for OCR functionality across different providers. - -This follows the same pattern as BaseLLMChatTest in tests/llm_translation/base_llm_unit_tests.py -""" - -import pytest -import litellm -import os -from abc import ABC, abstractmethod - - -# Test resources -TEST_IMAGE_PATH = "test_image_edit.png" -# Tiny in-repo PDF served via jsdelivr (sha-pinned, immutable). The arxiv -# PDF previously used here was several MB — once base64-encoded into the -# Vertex OCR request it ballooned cassettes past 100 MB per test. Keep -# the URL stable across runs so cassettes don't churn. -TEST_PDF_URL = ( - "https://cdn.jsdelivr.net/gh/BerriAI/litellm" - "@d769e81c90d453240c61fc572cdb27fae06a89d0" - "/tests/llm_translation/fixtures/dummy.pdf" -) - - -class BaseOCRTest(ABC): - """ - Abstract base test class that enforces common OCR tests across all providers. - - Each provider-specific test class should inherit from this and implement - get_base_ocr_call_args() to return provider-specific configuration. - """ - - @abstractmethod - def get_base_ocr_call_args(self) -> dict: - """Must return the base OCR call args for the specific provider""" - pass - - @pytest.mark.parametrize("sync_mode", [True, False]) - @pytest.mark.asyncio - async def test_basic_ocr_with_url(self, sync_mode): - """ - Test basic OCR with a public URL. - """ - litellm._turn_on_debug() - base_ocr_call_args = self.get_base_ocr_call_args() - print("BASE OCR Call args=", base_ocr_call_args) - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - try: - if sync_mode: - response = litellm.ocr( - document={"type": "document_url", "document_url": TEST_PDF_URL}, - **base_ocr_call_args, - ) - else: - response = await litellm.aocr( - document={"type": "document_url", "document_url": TEST_PDF_URL}, - **base_ocr_call_args, - ) - - print(f"\n{'='*80}") - print(f"Sync Mode: {sync_mode}") - print(f"Response type: {type(response)}") - print( - f"Response object: {response.object if hasattr(response, 'object') else 'N/A'}" - ) - - # Check if response has expected OCR format - assert hasattr(response, "pages"), "Response should have 'pages' attribute" - assert hasattr(response, "model"), "Response should have 'model' attribute" - assert hasattr( - response, "object" - ), "Response should have 'object' attribute" - assert ( - response.object == "ocr" - ), f"Expected object='ocr', got '{response.object}'" - - # Validate pages structure - assert isinstance(response.pages, list), "pages should be a list" - assert len(response.pages) > 0, "Should have at least one page" - - # Check first page structure - first_page = response.pages[0] - assert hasattr(first_page, "index"), "Page should have 'index' attribute" - assert hasattr( - first_page, "markdown" - ), "Page should have 'markdown' attribute" - - # Extract text from all pages for validation - total_text = "\n\n".join( - page.markdown for page in response.pages if page.markdown - ) - print(f"Total pages: {len(response.pages)}") - print(f"Total extracted text length: {len(total_text)} characters") - print(f"First 200 chars: {total_text[:200]}") - print(f"Model: {response.model}") - if response.usage_info: - print(f"Pages processed: {response.usage_info.pages_processed}") - print(f"{'='*80}\n") - - assert len(total_text) > 0, "Should extract some text from the document" - - ######################################################### - # validate we get a response cost in hidden parameters - ######################################################### - hidden_params = response._hidden_params - assert isinstance( - hidden_params, dict - ), "Hidden parameters should be a dictionary" - - print("response usage_info:", response.usage_info) - - response_cost = hidden_params.get("response_cost") - assert ( - response_cost is not None - ), "Response cost should be in hidden parameters" - assert response_cost > 0, "Response cost should be greater than 0" - print("response_cost=", response_cost) - - except litellm.RateLimitError as e: - error_msg = str(e) - if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg: - pytest.skip(f"Quota exceeded - {error_msg}") - else: - pytest.skip(f"Rate limit exceeded - {error_msg}") - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - except litellm.BadRequestError as e: - error_msg = str(e) - if ( - "URL_REJECTED" in error_msg - or "Cannot fetch content from the provided URL" in error_msg - ): - pytest.skip(f"URL rejected by provider - {error_msg}") - else: - pytest.fail(f"OCR call failed: {str(e)}") - except Exception as e: - pytest.fail(f"OCR call failed: {str(e)}") - - def test_ocr_response_structure(self): - """ - Test that the OCR response has the correct structure. - """ - litellm.set_verbose = True - base_ocr_call_args = self.get_base_ocr_call_args() - - try: - response = litellm.ocr( - document={"type": "document_url", "document_url": TEST_PDF_URL}, - **base_ocr_call_args, - ) - - # Validate response structure - assert hasattr(response, "pages"), "Response should have 'pages' attribute" - assert hasattr(response, "model"), "Response should have 'model' attribute" - assert hasattr( - response, "object" - ), "Response should have 'object' attribute" - assert hasattr( - response, "usage_info" - ), "Response should have 'usage_info' attribute" - - assert isinstance(response.pages, list), "pages should be a list" - assert len(response.pages) > 0, "Should have at least one page" - assert response.object == "ocr", "object should be 'ocr'" - - # Validate first page structure - first_page = response.pages[0] - assert hasattr(first_page, "index"), "Page should have 'index' attribute" - assert hasattr( - first_page, "markdown" - ), "Page should have 'markdown' attribute" - assert isinstance(first_page.markdown, str), "markdown should be a string" - - print(f"\nResponse structure validated:") - print(f" - object: {response.object}") - print(f" - model: {response.model}") - print(f" - pages: {len(response.pages)}") - if response.usage_info: - print(f" - pages_processed: {response.usage_info.pages_processed}") - print(f" - doc_size_bytes: {response.usage_info.doc_size_bytes}") - - except litellm.RateLimitError as e: - error_msg = str(e) - if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg: - pytest.skip(f"Quota exceeded - {error_msg}") - else: - pytest.skip(f"Rate limit exceeded - {error_msg}") - except litellm.InternalServerError: - pytest.skip("Model is overloaded") - except litellm.BadRequestError as e: - error_msg = str(e) - if ( - "URL_REJECTED" in error_msg - or "Cannot fetch content from the provided URL" in error_msg - ): - pytest.skip(f"URL rejected by provider - {error_msg}") - else: - pytest.fail(f"OCR response structure test failed: {str(e)}") - except Exception as e: - pytest.fail(f"OCR response structure test failed: {str(e)}") diff --git a/tests/ocr_tests/test_ocr_azure_ai.py b/tests/ocr_tests/test_ocr_azure_ai.py deleted file mode 100644 index acb44958fd9..00000000000 --- a/tests/ocr_tests/test_ocr_azure_ai.py +++ /dev/null @@ -1,29 +0,0 @@ -""" -Test OCR functionality with Azure AI API. - -Note: Azure AI OCR automatically converts URLs to base64 data URIs since -the Azure AI endpoint doesn't have internet access. -""" - -import os -from base_ocr_unit_tests import BaseOCRTest - - -class TestAzureAIOCR(BaseOCRTest): - """ - Test class for Azure AI OCR functionality. - Inherits from BaseOCRTest and provides Azure AI-specific configuration. - - Note: For Azure AI, LiteLLM will automatically convert URLs to base64 data URIs before - sending to the API, since Azure AI OCR endpoint doesn't have internet access. - """ - - def get_base_ocr_call_args(self) -> dict: - """ - Return the base OCR call args for Azure AI. - """ - return { - "model": "azure_ai/mistral-document-ai-2512", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - } diff --git a/tests/ocr_tests/test_ocr_azure_document_intelligence.py b/tests/ocr_tests/test_ocr_azure_document_intelligence.py index e6a2e5e5735..521f85ca8f7 100644 --- a/tests/ocr_tests/test_ocr_azure_document_intelligence.py +++ b/tests/ocr_tests/test_ocr_azure_document_intelligence.py @@ -1,53 +1,13 @@ -""" -Test OCR functionality with Azure Document Intelligence API. - -Azure Document Intelligence provides advanced document analysis capabilities -using the v4.0 (2024-11-30) API. -""" - -import os +"""Azure Document Intelligence request transformation: Mistral-shaped `pages` to Azure's query string.""" import pytest -from base_ocr_unit_tests import BaseOCRTest from litellm.constants import AZURE_DOCUMENT_INTELLIGENCE_API_VERSION from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( AzureDocumentIntelligenceOCRConfig, ) -class TestAzureDocumentIntelligenceOCR(BaseOCRTest): - """ - Test class for Azure Document Intelligence OCR functionality. - - Inherits from BaseOCRTest and provides Azure Document Intelligence-specific configuration. - - Tests the azure_ai/doc-intelligence/ provider route. - """ - - def get_base_ocr_call_args(self) -> dict: - """ - Return the base OCR call args for Azure Document Intelligence. - - Uses prebuilt-layout model which is closest to Mistral OCR format. - """ - # Check for required environment variables - api_key = os.environ.get("AZURE_DOCUMENT_INTELLIGENCE_API_KEY") - endpoint = os.environ.get("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") - - if not api_key or not endpoint: - pytest.skip( - "AZURE_DOCUMENT_INTELLIGENCE_API_KEY and AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT " - "environment variables are required for Azure Document Intelligence tests" - ) - - return { - "model": "azure_ai/doc-intelligence/prebuilt-layout", - "api_key": api_key, - "api_base": endpoint, - } - - class TestAzureDocumentIntelligencePagesParam: """ Unit tests for the Mistral-compatible `pages` parameter translation to @@ -101,7 +61,7 @@ class TestAzureDocumentIntelligencePagesParam: cfg.map_ocr_params({"pages": [True, False]}, {}, "prebuilt-layout") def test_map_ocr_params_unsupported_type_raises(self, cfg): - with pytest.raises(ValueError, match='based, Mistral-style\\) or a string like'): + with pytest.raises(ValueError, match="based, Mistral-style\\) or a string like"): cfg.map_ocr_params({"pages": 5}, {}, "prebuilt-layout") def test_get_complete_url_appends_pages_query(self, cfg): @@ -110,9 +70,7 @@ class TestAzureDocumentIntelligencePagesParam: model="azure_ai/doc-intelligence/prebuilt-layout", optional_params={"pages": "1-3,5"}, ) - assert ( - f"api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" in url - ), url + assert f"api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" in url, url assert "pages=1-3,5" in url, url assert "/documentintelligence/documentModels/prebuilt-layout:analyze" in url @@ -168,4 +126,3 @@ class TestAzureDocumentIntelligencePagesParam: assert "pages=3,4,5,6,7,8,9" in url assert req.data == {"urlSource": "https://example.com/x.pdf"} - diff --git a/tests/ocr_tests/test_ocr_matrix.py b/tests/ocr_tests/test_ocr_matrix.py new file mode 100644 index 00000000000..13cffbbc9a1 --- /dev/null +++ b/tests/ocr_tests/test_ocr_matrix.py @@ -0,0 +1,317 @@ +"""Live provider x auth x input coverage for ``litellm.ocr`` / ``litellm.aocr``. + +Each ``Case`` is one hand-picked cell, not the full cross product: every provider +exercises each of its credential kinds in both ``explicit`` (kwargs) and ``env`` +(monkeypatched environment) mode at least once, and every input kind a provider +accepts is exercised at least once. Sync and async are spread across the cells. +Every cell also checks the success callback saw the same response and cost. +""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import os +import re +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +import pytest + +import litellm +from litellm import Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.base_llm.ocr.transformation import OCRResponse + +Document = Mapping[str, object] +AuthMode = Literal["explicit", "env"] +CallStyle = Literal["sync", "async"] + + +@dataclass(frozen=True, slots=True) +class LoggedCall: + payload: Mapping[str, object] + response: object + + +class RecordingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger.__init__ is untyped + self.calls: Final[list[LoggedCall]] = [] # mutable-ok: append-only sink the callback hooks write into + + def _record(self, kwargs: Mapping[str, object], response_obj: object) -> None: + payload: Final = _string_keyed(kwargs.get("standard_logging_object")) + self.calls.append(LoggedCall(payload, response_obj)) + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self._record(kwargs, response_obj) + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self._record(kwargs, response_obj) + + async def wait_for_call(self, timeout: float = 10.0) -> LoggedCall: + deadline: Final = asyncio.get_running_loop().time() + timeout + while not self.calls: + assert asyncio.get_running_loop().time() < deadline, "success callback never fired" + await asyncio.sleep(0.05) + assert len(self.calls) == 1, self.calls + return self.calls[0] + + +@pytest.fixture +def logger(monkeypatch: pytest.MonkeyPatch) -> RecordingLogger: + recorder: Final = RecordingLogger() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + for registry in ("success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, registry, []) + return recorder + + +TESTS_DIR: Final = Path(__file__).resolve().parents[1] +PDF_PATH: Final = TESTS_DIR / "llm_translation" / "fixtures" / "dummy.pdf" +PNG_PATH: Final = TESTS_DIR / "image_gen_tests" / "test_image.png" +PINNED_CDN: Final = "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0" +PDF_URL: Final = f"{PINNED_CDN}/tests/llm_translation/fixtures/dummy.pdf" +PNG_URL: Final = f"{PINNED_CDN}/tests/image_gen_tests/test_image.png" +PDF_TEXT: Final = "Test PDF File" +PNG_TEXT: Final = "LiteLLM" + + +class _NamedReader(io.BytesIO): + def __init__(self, path: Path) -> None: + super().__init__(path.read_bytes()) + self.name: Final = path.name + + +def _data_uri(path: Path, mime: str) -> str: + return f"data:{mime};base64,{base64.b64encode(path.read_bytes()).decode()}" + + +@dataclass(frozen=True, slots=True) +class Input: + id: str + build: Callable[[], Document] + expected_text: str + + +PDF_BY_URL: Final = Input("pdf_url", lambda: {"type": "document_url", "document_url": PDF_URL}, PDF_TEXT) +PNG_BY_URL: Final = Input("image_url", lambda: {"type": "image_url", "image_url": PNG_URL}, PNG_TEXT) +PDF_DATA_URI: Final = Input( + "pdf_data_uri", + lambda: {"type": "document_url", "document_url": _data_uri(PDF_PATH, "application/pdf")}, + PDF_TEXT, +) +PNG_DATA_URI: Final = Input( + "image_data_uri", lambda: {"type": "image_url", "image_url": _data_uri(PNG_PATH, "image/png")}, PNG_TEXT +) +PDF_AS_PATH: Final = Input("pdf_path", lambda: {"type": "file", "file": PDF_PATH}, PDF_TEXT) +PDF_AS_BYTES: Final = Input( + "pdf_bytes", lambda: {"type": "file", "file": PDF_PATH.read_bytes(), "mime_type": "application/pdf"}, PDF_TEXT +) +PNG_AS_BYTES: Final = Input( + "image_bytes", lambda: {"type": "file", "file": PNG_PATH.read_bytes(), "mime_type": "image/png"}, PNG_TEXT +) +PNG_AS_FILE_OBJECT: Final = Input( + "image_file_object", lambda: {"type": "file", "file": _NamedReader(PNG_PATH)}, PNG_TEXT +) + + +@dataclass(frozen=True, slots=True) +class Secret: + """One credential value: the ``litellm.ocr`` kwarg it travels in, the env var litellm reads + when the kwarg is omitted, and the env var that holds the value in the test process.""" + + kwarg: str + env: str + source: str | None = None + + @property + def source_env(self) -> str: + return self.source or self.env + + +@dataclass(frozen=True, slots=True) +class Credential: + id: str + secrets: tuple[Secret, ...] + + +@dataclass(frozen=True, slots=True) +class Provider: + id: str + model: str + credentials: tuple[Credential, ...] + params: Mapping[str, str] = MappingProxyType({}) + + @property + def env_vars(self) -> frozenset[str]: + return frozenset(secret.env for credential in self.credentials for secret in credential.secrets) + + +MISTRAL_KEY: Final = Credential("api_key", (Secret("api_key", "MISTRAL_API_KEY"),)) +COHERE_KEY: Final = Credential("api_key", (Secret("api_key", "COHERE_API_KEY"),)) +REDUCTO_KEY: Final = Credential("api_key", (Secret("api_key", "REDUCTO_API_KEY"),)) + +AZURE_ENTRA_SECRETS: Final = ( + Secret("tenant_id", "AZURE_TENANT_ID", "AZURE_FOUNDRY_TENANT_ID"), + Secret("client_id", "AZURE_CLIENT_ID", "AZURE_FOUNDRY_ADMIN_CLIENT_ID"), + Secret("client_secret", "AZURE_CLIENT_SECRET", "AZURE_FOUNDRY_ADMIN_CLIENT_SECRET"), +) +AZURE_AI_BASE: Final = Secret("api_base", "AZURE_AI_API_BASE") +AZURE_AI_KEY: Final = Credential("api_key", (AZURE_AI_BASE, Secret("api_key", "AZURE_AI_API_KEY"))) +AZURE_AI_ENTRA: Final = Credential("entra", (AZURE_AI_BASE, *AZURE_ENTRA_SECRETS)) + +AZURE_DI_BASE: Final = Secret("api_base", "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") +AZURE_DI_KEY: Final = Credential("api_key", (AZURE_DI_BASE, Secret("api_key", "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"))) +AZURE_DI_ENTRA: Final = Credential("entra", (AZURE_DI_BASE, *AZURE_ENTRA_SECRETS)) + +VERTEX_SERVICE_ACCOUNT: Final = Credential( + "service_account", + (Secret("vertex_credentials", "VERTEXAI_CREDENTIALS"), Secret("vertex_project", "VERTEXAI_PROJECT")), +) + +MISTRAL: Final = Provider("mistral", "mistral/mistral-ocr-latest", (MISTRAL_KEY,)) +AZURE_AI_MISTRAL: Final = Provider( + "azure_ai_mistral", "azure_ai/mistral-document-ai-2512", (AZURE_AI_KEY, AZURE_AI_ENTRA) +) +AZURE_DOC_INTELLIGENCE: Final = Provider( + "azure_doc_intelligence", "azure_ai/doc-intelligence/prebuilt-layout", (AZURE_DI_KEY, AZURE_DI_ENTRA) +) +COHERE: Final = Provider("cohere", "cohere/parse-v5.0", (COHERE_KEY,)) +REDUCTO_V3: Final = Provider("reducto_v3", "reducto/parse-v3", (REDUCTO_KEY,)) +REDUCTO_LEGACY: Final = Provider("reducto_legacy", "reducto/parse-legacy", (REDUCTO_KEY,)) +VERTEX_MISTRAL: Final = Provider( + "vertex_mistral", + "vertex_ai/mistral-ocr-2505", + (VERTEX_SERVICE_ACCOUNT,), + MappingProxyType({"vertex_location": "us-central1"}), +) + + +@dataclass(frozen=True, slots=True) +class Case: + provider: Provider + credential: Credential + auth: AuthMode + document: Input + call: CallStyle + + @property + def id(self) -> str: + return f"{self.provider.id}-{self.credential.id}-{self.auth}-{self.document.id}-{self.call}" + + def bind_credentials(self, monkeypatch: pytest.MonkeyPatch) -> Mapping[str, str]: + """Clear every env var the provider could fall back to, then supply this case's values via kwargs or env.""" + values: Final = {secret: os.environ.get(secret.source_env) for secret in self.credential.secrets} + missing: Final = tuple(secret.source_env for secret, value in values.items() if not value) + if missing: + pytest.skip(f"{', '.join(missing)} not set") + for env_var in self.provider.env_vars: + monkeypatch.delenv(env_var, raising=False) + if self.auth == "explicit": + return {secret.kwarg: value for secret, value in values.items() if value} + for secret, value in values.items(): + monkeypatch.setenv(secret.env, value or "") + return {} + + async def run(self, credentials: Mapping[str, str]) -> OCRResponse: + kwargs: Final = {**self.provider.params, **credentials} + document: Final = self.document.build() + response: Final = ( + await litellm.aocr(model=self.provider.model, document=document, **kwargs) # pyright: ignore[reportUnknownMemberType] # @client erases the signature + if self.call == "async" + else litellm.ocr(model=self.provider.model, document=document, **kwargs) + ) + assert isinstance(response, OCRResponse) + return response + + +CASES: Final = ( + Case(MISTRAL, MISTRAL_KEY, "explicit", PDF_BY_URL, "sync"), + Case(MISTRAL, MISTRAL_KEY, "env", PNG_BY_URL, "async"), + Case(MISTRAL, MISTRAL_KEY, "explicit", PDF_AS_PATH, "sync"), + Case(MISTRAL, MISTRAL_KEY, "explicit", PNG_AS_BYTES, "async"), + Case(MISTRAL, MISTRAL_KEY, "explicit", PNG_AS_FILE_OBJECT, "sync"), + Case(AZURE_AI_MISTRAL, AZURE_AI_KEY, "explicit", PDF_BY_URL, "sync"), + Case(AZURE_AI_MISTRAL, AZURE_AI_KEY, "env", PNG_BY_URL, "async"), + Case(AZURE_AI_MISTRAL, AZURE_AI_ENTRA, "explicit", PDF_AS_PATH, "sync"), + Case(AZURE_AI_MISTRAL, AZURE_AI_ENTRA, "env", PDF_DATA_URI, "async"), + Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_KEY, "explicit", PDF_BY_URL, "sync"), + Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_KEY, "env", PNG_AS_BYTES, "async"), + Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_ENTRA, "explicit", PNG_BY_URL, "async"), + Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_ENTRA, "env", PDF_AS_PATH, "sync"), + Case(COHERE, COHERE_KEY, "explicit", PNG_BY_URL, "sync"), + Case(COHERE, COHERE_KEY, "env", PNG_DATA_URI, "async"), + Case(REDUCTO_V3, REDUCTO_KEY, "explicit", PDF_AS_PATH, "sync"), + Case(REDUCTO_V3, REDUCTO_KEY, "env", PNG_AS_BYTES, "async"), + Case(REDUCTO_V3, REDUCTO_KEY, "explicit", PDF_DATA_URI, "async"), + Case(REDUCTO_LEGACY, REDUCTO_KEY, "explicit", PDF_AS_BYTES, "sync"), + Case(VERTEX_MISTRAL, VERTEX_SERVICE_ACCOUNT, "explicit", PDF_BY_URL, "sync"), + Case(VERTEX_MISTRAL, VERTEX_SERVICE_ACCOUNT, "env", PNG_BY_URL, "async"), +) + + +def _response_cost(response: OCRResponse) -> float: + response_cost: Final[object] = response._hidden_params.get("response_cost") # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType] # response_cost is only surfaced on _hidden_params + assert isinstance(response_cost, float) and response_cost > 0 + return response_cost + + +def _assert_ocr_response(response: OCRResponse, model: str, expected_text: str) -> None: + assert response.object == "ocr" + assert response.model == model.split("/", 1)[1] + assert [page.index for page in response.pages] == list(range(len(response.pages))) + text: Final = re.sub(r"\s+", " ", " ".join(page.markdown for page in response.pages)) + assert expected_text.lower() in text.lower(), text + assert response.usage_info is not None + assert response.usage_info.pages_processed == len(response.pages) + _response_cost(response) + + +def _string_keyed(value: object) -> Mapping[str, object]: + assert isinstance(value, Mapping), type(value) + items: Final = tuple(value.items()) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportUnknownArgumentType] # narrowed from object + return MappingProxyType({str(key): value for key, value in items}) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # narrowed from object + + +def _assert_logged(logged: LoggedCall, response: OCRResponse, model: str, logged_model: str, call: CallStyle) -> None: + assert isinstance(logged.response, OCRResponse) + assert logged.response.pages == response.pages + assert logged.payload["status"] == "success" + assert logged.payload["call_type"] == ("aocr" if call == "async" else "ocr") + assert logged.payload["custom_llm_provider"] == model.split("/", 1)[0] + assert logged.payload["model"] == logged_model + assert logged.payload["response_cost"] == _response_cost(response) + + +@pytest.mark.parametrize("case", CASES, ids=[case.id for case in CASES]) +async def test_ocr(case: Case, monkeypatch: pytest.MonkeyPatch, logger: RecordingLogger) -> None: + credentials: Final = case.bind_credentials(monkeypatch) + response: Final = await case.run(credentials) + _assert_ocr_response(response, case.provider.model, case.document.expected_text) + _assert_logged(await logger.wait_for_call(), response, case.provider.model, response.model, case.call) + + +async def test_router_aocr(monkeypatch: pytest.MonkeyPatch, logger: RecordingLogger) -> None: + case: Final = Case(MISTRAL, MISTRAL_KEY, "explicit", PDF_BY_URL, "async") + router: Final = Router( + model_list=[ + { + "model_name": "ocr-alias", + "litellm_params": {"model": MISTRAL.model, **case.bind_credentials(monkeypatch)}, + } + ] + ) + response: Final = await router.aocr(model="ocr-alias", document=PDF_BY_URL.build()) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # Router.aocr is untyped + assert isinstance(response, OCRResponse) + _assert_ocr_response(response, MISTRAL.model, PDF_TEXT) + _assert_logged(await logger.wait_for_call(), response, MISTRAL.model, MISTRAL.model, case.call) diff --git a/tests/ocr_tests/test_ocr_mistral.py b/tests/ocr_tests/test_ocr_mistral.py deleted file mode 100644 index cdc093620b2..00000000000 --- a/tests/ocr_tests/test_ocr_mistral.py +++ /dev/null @@ -1,92 +0,0 @@ -""" -Test OCR functionality with Mistral API. -""" - -import os -import sys -import pytest -import litellm -from litellm import Router -from base_ocr_unit_tests import BaseOCRTest, TEST_PDF_URL - - -class TestMistralOCR(BaseOCRTest): - """ - Test class for Mistral OCR functionality. - """ - - def get_base_ocr_call_args(self) -> dict: - """Return the base OCR call args for Mistral""" - return { - "model": "mistral/mistral-ocr-latest", - "api_key": os.getenv("MISTRAL_API_KEY"), - } - - -@pytest.mark.asyncio -async def test_router_aocr_with_mistral(): - """ - Test OCR with Router using Mistral OCR deployment. - """ - litellm.set_verbose = True - - # Create router with Mistral OCR deployment - router = Router( - model_list=[ - { - "model_name": "mistral-ocr", - "litellm_params": { - "model": "mistral/mistral-ocr-latest", - "api_key": os.getenv("MISTRAL_API_KEY"), - }, - } - ] - ) - - try: - # Call OCR through router - response = await router.aocr( - model="mistral-ocr", - document={"type": "document_url", "document_url": TEST_PDF_URL}, - ) - - print(f"\n{'='*80}") - print("Router OCR Test") - print(f"Response type: {type(response)}") - print( - f"Response object: {response.object if hasattr(response, 'object') else 'N/A'}" - ) - - # Check if response has expected Mistral OCR format - assert hasattr(response, "pages"), "Response should have 'pages' attribute" - assert hasattr(response, "model"), "Response should have 'model' attribute" - assert hasattr(response, "object"), "Response should have 'object' attribute" - assert ( - response.object == "ocr" - ), f"Expected object='ocr', got '{response.object}'" - - # Validate pages structure - assert isinstance(response.pages, list), "pages should be a list" - assert len(response.pages) > 0, "Should have at least one page" - - # Check first page structure - first_page = response.pages[0] - assert hasattr(first_page, "index"), "Page should have 'index' attribute" - assert hasattr(first_page, "markdown"), "Page should have 'markdown' attribute" - - # Extract text from all pages for validation - total_text = "\n\n".join( - page.markdown for page in response.pages if page.markdown - ) - print(f"Total pages: {len(response.pages)}") - print(f"Total extracted text length: {len(total_text)} characters") - print(f"First 200 chars: {total_text[:200]}") - print(f"Model: {response.model}") - if response.usage_info: - print(f"Pages processed: {response.usage_info.pages_processed}") - print(f"{'='*80}\n") - - assert len(total_text) > 0, "Should extract some text from the document" - - except Exception as e: - pytest.fail(f"Router OCR call failed: {str(e)}") diff --git a/tests/ocr_tests/test_ocr_vertex_ai.py b/tests/ocr_tests/test_ocr_vertex_ai.py index 1842eb063a5..beddc9cd35e 100644 --- a/tests/ocr_tests/test_ocr_vertex_ai.py +++ b/tests/ocr_tests/test_ocr_vertex_ai.py @@ -1,117 +1,8 @@ -""" -Test OCR functionality with Vertex AI OCR APIs (Mistral and DeepSeek). +"""Vertex AI OCR config routing and DeepSeek request shaping (no network).""" -Note: Vertex AI OCR automatically converts URLs to base64 data URIs since -the Vertex AI endpoint doesn't have internet access. -""" - -import json -import os -import tempfile from typing import Final import pytest -from base_ocr_unit_tests import BaseOCRTest - - -def load_vertex_ai_credentials(): - """Load Vertex AI credentials for tests""" - # Define the path to the vertex_key.json file - print("loading vertex ai credentials") - filepath = os.path.dirname(os.path.abspath(__file__)) - vertex_key_path = filepath + "/vertex_key.json" - - # Read the existing content of the file or create an empty dictionary - try: - with open(vertex_key_path, "r") as file: - # Read the file content - print("Read vertexai file path") - content = file.read() - - # If the file is empty or not valid JSON, create an empty dictionary - if not content or not content.strip(): - service_account_key_data = {} - else: - # Attempt to load the existing JSON content - file.seek(0) - service_account_key_data = json.load(file) - except FileNotFoundError: - # If the file doesn't exist, create an empty dictionary - service_account_key_data = {} - - # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - private_key = private_key.replace("\\n", "\n") - service_account_key_data["private_key_id"] = private_key_id - service_account_key_data["private_key"] = private_key - - # Create a temporary file - with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: - # Write the updated content to the temporary files - json.dump(service_account_key_data, temp_file, indent=2) - - # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) - - -class TestVertexAIMistralOCR(BaseOCRTest): - """ - Test class for Vertex AI Mistral OCR functionality. - Inherits from BaseOCRTest and provides Vertex AI-specific configuration. - - Note: For Vertex AI, LiteLLM will automatically convert URLs to base64 data URIs before - sending to the API, since Vertex AI OCR endpoint doesn't have internet access. - """ - - def setup_method(self): - if os.environ.get("LITELLM_RUN_LIVE_VERTEX_MISTRAL_OCR_TESTS") != "1": - pytest.skip("Live Vertex AI Mistral OCR E2E tests are opt-in") - if os.environ.get("CASSETTE_REDIS_URL"): - pytest.skip( - "Live Vertex AI Mistral OCR E2E tests cannot run under VCR replay" - ) - - def get_base_ocr_call_args(self) -> dict: - """ - Return the base OCR call args for Vertex AI Mistral OCR. - """ - load_vertex_ai_credentials() - return { - "model": "vertex_ai/mistral-ocr-2505", - "vertex_location": "us-central1", - } - - -class TestVertexAIDeepSeekOCR(BaseOCRTest): - """ - Test class for Vertex AI DeepSeek OCR functionality. - Inherits from BaseOCRTest and provides Vertex AI-specific configuration. - - Note: DeepSeek OCR uses the chat completion API format through the openapi endpoint. - Note: DeepSeek OCR does not support PDF URLs - only image URLs and base64 data. - """ - - def get_base_ocr_call_args(self) -> dict: - """ - Return the base OCR call args for Vertex AI DeepSeek OCR. - """ - load_vertex_ai_credentials() - return { - "model": "vertex_ai/deepseek-ocr-maas", - "vertex_location": "us-central1", - } - - # Skip PDF URL tests for DeepSeek OCR as it doesn't support PDF URLs - @pytest.mark.skip(reason="DeepSeek OCR does not support PDF URLs") - async def test_basic_ocr_with_url(self, sync_mode): - """Skip this test for DeepSeek OCR - PDF URLs not supported""" - pass - - @pytest.mark.skip(reason="DeepSeek OCR does not support PDF URLs") - def test_ocr_response_structure(self): - """Skip this test for DeepSeek OCR - PDF URLs not supported""" - pass def test_vertex_ai_ocr_routing(): @@ -126,21 +17,19 @@ def test_vertex_ai_ocr_routing(): # Test DeepSeek OCR routing deepseek_config = get_vertex_ai_ocr_config("vertex_ai/deepseek-ocr-maas") - assert isinstance( - deepseek_config, VertexAIDeepSeekOCRConfig - ), "DeepSeek model should route to VertexAIDeepSeekOCRConfig" + assert isinstance(deepseek_config, VertexAIDeepSeekOCRConfig), ( + "DeepSeek model should route to VertexAIDeepSeekOCRConfig" + ) # Test Mistral OCR routing (should use default VertexAIOCRConfig) mistral_config = get_vertex_ai_ocr_config("vertex_ai/mistral-ocr-2505") - assert isinstance( - mistral_config, VertexAIOCRConfig - ), "Mistral model should route to VertexAIOCRConfig" + assert isinstance(mistral_config, VertexAIOCRConfig), "Mistral model should route to VertexAIOCRConfig" # Test other DeepSeek variants deepseek_variant = get_vertex_ai_ocr_config("vertex_ai/deepseek-ocr-maas") - assert isinstance( - deepseek_variant, VertexAIDeepSeekOCRConfig - ), "DeepSeek variant should route to VertexAIDeepSeekOCRConfig" + assert isinstance(deepseek_variant, VertexAIDeepSeekOCRConfig), ( + "DeepSeek variant should route to VertexAIDeepSeekOCRConfig" + ) @pytest.mark.parametrize("model", ("deepseek-ocr-maas", "deepseek-ai/deepseek-ocr-maas")) diff --git a/tests/ocr_tests/vertex_key.json b/tests/ocr_tests/vertex_key.json deleted file mode 100644 index 800969fb305..00000000000 --- a/tests/ocr_tests/vertex_key.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "type": "service_account", - "project_id": "litellm-ci-cd", - "private_key_id": "", - "private_key": "", - "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", - "client_id": "116563532503305622785", - "auth_uri": "https://accounts.google.com/o/oauth2/auth", - "token_uri": "https://oauth2.googleapis.com/token", - "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", - "universe_domain": "googleapis.com" -} From 0b5b69ea3ae4aa2c8aeb5764240046c85713d5a0 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 21:29:29 +0000 Subject: [PATCH 077/102] fix(deepgram): forward only the first model and language values to /listen Authorization and pricing read the first model and language query value, but the raw query was forwarded, so Deepgram (which honours the last repeated value) could be sent a model the key was never allowed. Later duplicates of those two keys are now dropped before the upstream URL is built Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/deepgram/common_utils.py | 21 ++++++++--- .../deepgram/test_deepgram_common_utils.py | 35 +++++++++++++++++-- .../test_deepgram_ws_passthrough_routes.py | 30 ++++++++++++++++ 3 files changed, 78 insertions(+), 8 deletions(-) diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index 676391dc744..9b071ac8321 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -26,6 +26,7 @@ DEEPGRAM_LISTEN_ADDON_PRICING_PARAMS: Final = MappingProxyType( } ) _DISABLED_PARAM_VALUES: Final = frozenset({"", "false"}) +_SINGLE_VALUED_PARAMS: Final = frozenset({"model", "language"}) class DeepgramException(BaseLLMException): @@ -36,13 +37,23 @@ def deepgram_listen_requested_model(query_string: str) -> str: return httpx.QueryParams(query_string).get("model") or DEEPGRAM_LISTEN_DEFAULT_MODEL +def _first_occurrences(query_string: str) -> httpx.QueryParams: + """Authorization and pricing read the first ``model`` and ``language`` value; Deepgram must not see a second one.""" + items: Final = httpx.QueryParams(query_string).multi_items() + return httpx.QueryParams( + tuple( + (key, value) + for index, (key, value) in enumerate(items) + if key not in _SINGLE_VALUED_PARAMS or all(earlier != key for earlier, _ in items[:index]) + ) + ) + + def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str: listen_url: Final = httpx.URL(f"{(api_base or DEEPGRAM_DEFAULT_API_BASE).rstrip('/')}/listen") websocket_url: Final = listen_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(listen_url.scheme, listen_url.scheme)) - params: Final = httpx.QueryParams(query_string) - query: Final = ( - query_string if params.get("model") else str(params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL)) - ) + params: Final = _first_occurrences(query_string) + query: Final = params if params.get("model") else params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL) return f"{websocket_url}?{query}" @@ -64,7 +75,7 @@ def deepgram_listen_pricing_model(upstream_url: str) -> str: the multilingual streaming entry when ``language=multi``, otherwise the model's own streaming entry. Pre-recorded entries are never a substitute: Deepgram prices the two products differently.""" streaming: Final = f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{deepgram_listen_model(upstream_url)}" - language: Final = parse_qs(urlparse(upstream_url).query).get("language", ("",))[-1] + language: Final = parse_qs(urlparse(upstream_url).query).get("language", ("",))[0] if language.strip().lower() == DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE: return f"{streaming}{DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX}" return streaming diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py index a1fa8f26b70..530888b70c4 100644 --- a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py @@ -1,6 +1,7 @@ import math from collections.abc import Mapping, Sequence from typing import Final +from urllib.parse import parse_qs, urlparse import pytest @@ -76,6 +77,24 @@ def _metadata(duration: object, channels: object = 1) -> dict[str, object]: "wss://dg.internal/v1/listen?model=nova-3&keywords=a&keywords=b", id="repeated keys preserved", ), + pytest.param( + None, + "model=nova-2&encoding=linear16&model=nova-3", + "wss://api.deepgram.com/v1/listen?model=nova-2&encoding=linear16", + id="only the authorized first model reaches deepgram", + ), + pytest.param( + None, + "language=en&model=nova-3&language=multi", + "wss://api.deepgram.com/v1/listen?language=en&model=nova-3", + id="only the priced first language reaches deepgram", + ), + pytest.param( + None, + "model=&model=nova-2", + "wss://api.deepgram.com/v1/listen?model=nova-3", + id="blank first model is the default, later models dropped", + ), ], ) def test_deepgram_listen_websocket_target(api_base: str | None, query_string: str, expected: str): @@ -210,12 +229,22 @@ def test_deepgram_listen_model_comes_from_the_upstream_query(upstream_url: str, @pytest.mark.parametrize( "query_string", - ["model=nova-2&language=en", "language=en", "model=&language=en", "", "model=nova-3-medical"], + [ + "model=nova-2&language=en", + "language=en", + "model=&language=en", + "", + "model=nova-3-medical", + "model=nova-2&model=nova-3", + "model=&model=nova-3-medical", + ], ) -def test_requested_model_is_the_model_the_upstream_target_will_carry(query_string: str): +def test_requested_model_is_the_only_model_the_upstream_target_carries(query_string: str): """Authorization runs against ``deepgram_listen_requested_model``; the upstream URL is built separately, so the - two must always agree or a key could be authorized for one model and reach another.""" + two must always agree or a key could be authorized for one model and reach another. Deepgram reads the last + repeated ``model``, so the target must carry exactly one.""" target: Final = deepgram_listen_websocket_target(None, query_string) + assert parse_qs(urlparse(target).query)["model"] == [deepgram_listen_requested_model(query_string)] assert deepgram_listen_requested_model(query_string) == deepgram_listen_model(target) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 4eb183b14ce..44533f35c72 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -421,6 +421,36 @@ def test_deepgram_listen_authorizes_the_model_it_will_actually_send_upstream(que assert relay.calls == [] +def test_deepgram_listen_strips_a_second_model_that_would_outrank_the_authorized_one(monkeypatch): + """Deepgram honours the last repeated ``model``; auth and pricing read the first. A key allowed only ``nova-2`` + must not smuggle ``nova-3`` past authorization behind an authorized first value.""" + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setattr(litellm, "max_budget", 0.0) + _price_nova_2_streaming(monkeypatch) + cache = asyncio.run(_cache_restricted_key("sk-only-nova-2", ["nova-2"])) + relay = _FakeRelay() + client = TestClient(_app_with_relay(relay)) + + with ( + patch(GET_CREDENTIALS, return_value="dg-provider-key"), + patch.multiple( # test-quality-ok: the real key auth path reads these proxy_server globals and has no injection seam + "litellm.proxy.proxy_server", + master_key="sk-master", + prisma_client=MagicMock(), + user_api_key_cache=cache, + llm_model_list=None, + llm_router=None, + ), + ): + with client.websocket_connect( + "/deepgram/v1/listen?model=nova-2&language=en&model=nova-3&language=multi", + headers={"Authorization": "Bearer sk-only-nova-2"}, + ): + pass + + assert [call.target for call in relay.calls] == ["wss://api.deepgram.com/v1/listen?model=nova-2&language=en"] + + def test_deepgram_listen_echoes_the_browser_subprotocol_that_carries_the_litellm_key(monkeypatch): """Browsers cannot set headers, so they send the key as a subprotocol and abort the handshake unless the server echoes that subprotocol back; the key itself must still stay off the upstream connection.""" From f72b7155acf0b17e61f10733fb897f1f5defb832 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 21:40:13 +0000 Subject: [PATCH 078/102] fix(ocr): map Rust upstream 401/403 to the public auth exceptions The httpx.Response built for a Rust upstream failure had no request attached, so constructing openai.AuthenticationError raised RuntimeError inside the exception mapper and every bad-key OCR call surfaced as APIConnectionError 500 instead of AuthenticationError 401 (the Python path already returned 401) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/rust_bridge/ocr/route_host.py | 10 +++++++--- tests/test_litellm/rust_bridge/ocr/test_route_host.py | 11 +++++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/litellm/rust_bridge/ocr/route_host.py b/litellm/rust_bridge/ocr/route_host.py index 0bc7b383eea..277fdceb734 100644 --- a/litellm/rust_bridge/ocr/route_host.py +++ b/litellm/rust_bridge/ocr/route_host.py @@ -27,13 +27,17 @@ class UpstreamFailure(Exception): self.__cause__ = cause -def _upstream_failure(error: Exception) -> Exception: +def _upstream_failure(error: Exception, request: LiteLLMOcrRequest) -> Exception: try: status, body = _UPSTREAM_ARGS.validate_python(error.args) headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) except ValidationError: return error - return UpstreamFailure(httpx.Response(status, content=body.encode(), headers=headers), error) + http_request: Final = httpx.Request("POST", request.api_base or "https://docs.litellm.ai/docs") + return UpstreamFailure( + httpx.Response(status, content=body.encode(), headers=headers, request=http_request), + error, + ) def response(value: Mapping[str, object]) -> OCRResponse: @@ -57,7 +61,7 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: model=request.model.removeprefix(f"{request_provider}/"), llm_provider=request_provider, ) - original: Final = _upstream_failure(error) + original: Final = _upstream_failure(error, request) public_error: Final = failures.map_failure(original, request.model, request_provider, arguments(request)) if isinstance(original, UpstreamFailure) and public_error.__context__ is original: public_error.__context__ = error diff --git a/tests/test_litellm/rust_bridge/ocr/test_route_host.py b/tests/test_litellm/rust_bridge/ocr/test_route_host.py index a328579400c..699492e4424 100644 --- a/tests/test_litellm/rust_bridge/ocr/test_route_host.py +++ b/tests/test_litellm/rust_bridge/ocr/test_route_host.py @@ -59,6 +59,17 @@ def test_map_failure_builds_public_error_from_upstream_status_and_headers() -> N assert public_error.llm_provider == "mistral" +def test_map_failure_maps_upstream_401_to_authentication_error() -> None: + error: Final = RustUpstreamError(401, '{"message": "Unauthorized"}', ()) + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert isinstance(public_error, litellm.AuthenticationError) + assert public_error.status_code == 401 + assert public_error.response.text == '{"message": "Unauthorized"}' + assert public_error.__context__ is error + + def test_map_failure_leaves_non_upstream_errors_unwrapped() -> None: error: Final = RuntimeError("bridge exploded") From cf05466a27ad9409d1edcd9697ddb826ca8f93da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:54:18 -0700 Subject: [PATCH 079/102] fix(gemini): map every documented finishReason and reset per-candidate state A content-less candidate is now kept as a choice whenever it carries a finishReason, with the raw value on the choice's provider_specific_fields. NO_IMAGE, IMAGE_RECITATION, IMAGE_OTHER and ESCALATION map to content_filter; UNEXPECTED_TOOL_CALL and MISSING_THOUGHT_SIGNATURE map to stop. The /v1/responses bridge reports content_filter and refusal as incomplete with incomplete_details, and tool calls and reasoning no longer leak from one candidate into the next. --- litellm/litellm_core_utils/core_helpers.py | 5 ++ .../vertex_and_google_ai_studio_gemini.py | 29 ++++---- .../transformation.py | 33 ++++++--- .../litellm_core_utils/test_core_helpers.py | 6 ++ ...test_vertex_and_google_ai_studio_gemini.py | 68 +++++++++++++++++++ .../test_litellm_completion_responses.py | 2 +- 6 files changed, 117 insertions(+), 26 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 2615c58de2c..d29b1fc74ef 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -225,6 +225,11 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { "TOO_MANY_TOOL_CALLS": "stop", "MALFORMED_RESPONSE": "stop", "NO_IMAGE": "content_filter", + "IMAGE_RECITATION": "content_filter", + "IMAGE_OTHER": "content_filter", + "ESCALATION": "content_filter", + "UNEXPECTED_TOOL_CALL": "stop", + "MISSING_THOUGHT_SIGNATURE": "stop", # Zhipu GLM "network_error": "stop", "sensitive": "content_filter", diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 61bac4793a8..a95b845718a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1348,6 +1348,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "TOO_MANY_TOOL_CALLS", "MALFORMED_RESPONSE", "NO_IMAGE", + "IMAGE_RECITATION", + "IMAGE_OTHER", + "ESCALATION", + "UNEXPECTED_TOOL_CALL", + "MISSING_THOUGHT_SIGNATURE", } ) @@ -2243,7 +2248,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): image_response: list[ImageURLListItem] | None = None chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} chat_completion_logprobs: ChoiceLogprobs | None = None - tools: list[ChatCompletionToolCallChunk] | None = [] + tools: list[ChatCompletionToolCallChunk] | None = None functions: ChatCompletionToolCallFunctionChunk | None = None thinking_blocks: list[ChatCompletionThinkingBlock] | None = None reasoning_content: str | None = None @@ -2358,11 +2363,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_invocation_fields["server_side_tool_invocations"] = server_side_tool_invocations chat_completion_message["provider_specific_fields"] = tool_invocation_fields - if candidate.get("finishReason"): - finish_reason_fields = chat_completion_message.get("provider_specific_fields") or {} - finish_reason_fields["native_finish_reason"] = candidate.get("finishReason") - chat_completion_message["provider_specific_fields"] = finish_reason_fields - if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( chat_completion_message=chat_completion_message, @@ -2375,15 +2375,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) model_response.choices.append(choice) elif isinstance(model_response, ModelResponse): + native_finish_reason = candidate.get("finishReason") choice = litellm.Choices( finish_reason=VertexGeminiConfig._check_finish_reason( - chat_completion_message, candidate.get("finishReason") + chat_completion_message, native_finish_reason ), index=candidate.get("index", idx), message=chat_completion_message, logprobs=chat_completion_logprobs, enhancements=None, - provider_specific_fields=chat_completion_message.get("provider_specific_fields"), + provider_specific_fields=( + {"native_finish_reason": native_finish_reason} if native_finish_reason is not None else None + ), ) model_response.choices.append(choice) @@ -3181,12 +3184,10 @@ class ModelResponseIterator: self.has_seen_tool_calls = True break - # _process_candidates skips candidates without a "content" part, so a - # content-less chunk leaves choices empty and the downstream streaming - # handler hits IndexError on choices[0]. This covers the final chunk - # (finishReason, no content) and mid-stream metadata-only chunks - # (grounding/web-search/thought, no content and no finishReason — seen - # with web_search + reasoning) by emitting an empty-delta choice. + # _process_candidates skips candidates with neither "content" nor + # "finishReason", so a metadata-only chunk (grounding/web-search/thought, + # seen with web_search + reasoning) leaves choices empty and the downstream + # streaming handler hits IndexError on choices[0]. Emit an empty-delta choice. if not model_response.choices and _candidates: from litellm.types.utils import Delta, StreamingChoices diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 7792bc11405..961fffd3d42 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -61,6 +61,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParamFunctionChunk, ChatCompletionUserMessage, GenericChatCompletionMessage, + IncompleteDetails, InputTokensDetails, OpenAIChatCompletionTextObject, OpenAIMcpServerTool, @@ -2295,6 +2296,21 @@ class LiteLLMCompletionResponsesConfig: # Default to completed for unknown finish reasons return "completed" + @staticmethod + def _incomplete_details_for_finish_reason( + finish_reason: str | None, + existing: IncompleteDetails | None, + ) -> IncompleteDetails | None: + if existing is not None: + return existing + match finish_reason: + case "length": + return IncompleteDetails(reason="max_output_tokens") + case "content_filter" | "refusal": + return IncompleteDetails(reason="content_filter") + case _: + return None + @staticmethod def _tool_call_id_from_responses_item(item_id: str | None, call_id: str | None) -> str: """Bedrock Mantle returns a non-unique, index-based ``call_id`` (``call_0``, @@ -2411,17 +2427,10 @@ class LiteLLMCompletionResponsesConfig: if choices and len(choices) > 0: finish_reason = choices[0].finish_reason - status: Final[ResponsesAPIStatus] = ( - LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status(finish_reason) + incomplete_details: Final = LiteLLMCompletionResponsesConfig._incomplete_details_for_finish_reason( + finish_reason=finish_reason, + existing=getattr(chat_completion_response, "incomplete_details", None), ) - incomplete_details = getattr(chat_completion_response, "incomplete_details", None) - if incomplete_details is None and status == "incomplete": - from openai.types.responses.response import IncompleteDetails - - if finish_reason == "length": - incomplete_details = IncompleteDetails(reason="max_output_tokens") - elif finish_reason in ["content_filter", "refusal"]: - incomplete_details = IncompleteDetails(reason="content_filter") responses_api_response: Final[ResponsesAPIResponse] = ResponsesAPIResponse( id=chat_completion_response.id, @@ -2447,7 +2456,9 @@ class LiteLLMCompletionResponsesConfig: max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), previous_response_id=getattr(chat_completion_response, "previous_response_id", None), reasoning=None, - status=status, + status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( + finish_reason + ), text={}, truncation=getattr(chat_completion_response, "truncation", None), usage=LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index e937be47441..b2ad13c205e 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -151,6 +151,12 @@ class TestMapFinishReasonGemini: ("IMAGE_PROHIBITED_CONTENT", "content_filter"), ("TOO_MANY_TOOL_CALLS", "stop"), ("MALFORMED_RESPONSE", "stop"), + ("NO_IMAGE", "content_filter"), + ("IMAGE_RECITATION", "content_filter"), + ("IMAGE_OTHER", "content_filter"), + ("ESCALATION", "content_filter"), + ("UNEXPECTED_TOOL_CALL", "stop"), + ("MISSING_THOUGHT_SIGNATURE", "stop"), ], ) def test_gemini_finish_reasons(self, gemini_reason, expected): diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 216d6f0db7e..4ee199bd9af 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -968,6 +968,12 @@ def test_finish_reason_unspecified_and_malformed_function_call(): # Test new Gemini finish reasons assert finish_reason_mappings["TOO_MANY_TOOL_CALLS"] == "stop" assert finish_reason_mappings["MALFORMED_RESPONSE"] == "stop" + assert finish_reason_mappings["NO_IMAGE"] == "content_filter" + assert finish_reason_mappings["IMAGE_RECITATION"] == "content_filter" + assert finish_reason_mappings["IMAGE_OTHER"] == "content_filter" + assert finish_reason_mappings["ESCALATION"] == "content_filter" + assert finish_reason_mappings["UNEXPECTED_TOOL_CALL"] == "stop" + assert finish_reason_mappings["MISSING_THOUGHT_SIGNATURE"] == "stop" def test_vertex_ai_usage_metadata_response_token_count(): @@ -6219,3 +6225,65 @@ def test_gemini_candidate_other_finish_reasons_no_content(): ) assert responses_length.status == "incomplete" assert responses_length.incomplete_details.reason == "max_output_tokens" + + +def test_gemini_candidate_with_finish_reason_no_content_streaming_chunk(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + chunk: Final = { + "candidates": [{"finishReason": "NO_IMAGE", "index": 0}], + "usageMetadata": {"promptTokenCount": 19, "candidatesTokenCount": 0, "totalTokenCount": 19}, + } + iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) + + streaming_chunk: Final = iterator.chunk_parser(chunk) + + assert len(streaming_chunk.choices) == 1 + assert streaming_chunk.choices[0].finish_reason == "content_filter" + assert streaming_chunk.choices[0].delta.content is None + assert streaming_chunk.choices[0].delta.tool_calls is None + + +def test_gemini_multi_candidate_messages_do_not_share_state(): + config: Final = VertexGeminiConfig() + completion_response: Final = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"text": "Let me check the weather.", "thought": True}, + {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}}, + ], + }, + "finishReason": "STOP", + "index": 0, + }, + { + "content": {"role": "model", "parts": [{"text": "It is sunny in Paris."}]}, + "finishReason": "STOP", + "index": 1, + }, + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 20, "totalTokenCount": 30}, + } + + resp: Final = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + assert len(resp.choices) == 2 + assert resp.choices[0].finish_reason == "tool_calls" + assert resp.choices[0].message.tool_calls[0].function.name == "get_weather" + assert resp.choices[0].message.reasoning_content == "Let me check the weather." + assert resp.choices[1].finish_reason == "stop" + assert resp.choices[1].message.content == "It is sunny in Paris." + assert resp.choices[1].message.tool_calls is None + assert getattr(resp.choices[1].message, "reasoning_content", None) is None + assert resp.choices[1].provider_specific_fields["native_finish_reason"] == "STOP" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 849c997838d..52d06acfd64 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4909,7 +4909,7 @@ class TestStreamingSnapshotItemIds: def test_transform_chat_completion_response_incomplete_details(): - from openai.types.responses.response import IncompleteDetails + from litellm.types.llms.openai import IncompleteDetails resp_length = ModelResponse( id="resp-length", From 55249c7128bca0c47641f561c30e96040e734d69 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:58:58 -0700 Subject: [PATCH 080/102] fix: set vertex gemma-4-26b-a4b-it-maas context window to 262144 --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- .../gemma/test_vertex_ai_gemma_global_endpoint.py | 12 +++++++++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 377e30a23b7..ef4fce37f23 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -50721,7 +50721,7 @@ "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 377e30a23b7..ef4fce37f23 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -50721,7 +50721,7 @@ "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py index e9b58622a4b..9d08daa0a81 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py @@ -30,7 +30,7 @@ from litellm.types.llms.vertex_ai import VertexPartnerProvider _GEMMA_MODEL_COST_ENTRY = { "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -180,6 +180,16 @@ class TestCreateVertexURLGemma: # --------------------------------------------------------------------------- +def test_gemma_maas_context_window_matches_google(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + info = litellm.get_model_info("vertex_ai/google/gemma-4-26b-a4b-it-maas") + + assert info["max_input_tokens"] == 262144 + assert info["max_output_tokens"] == 128000 + + # --------------------------------------------------------------------------- # Integration tests: verify payloads reach the global OpenAI endpoint # From 19ef8e47a6dc3453b414b2efe556b9e2893c3e24 Mon Sep 17 00:00:00 2001 From: joshua Date: Fri, 18 Sep 2026 22:01:56 +0000 Subject: [PATCH 081/102] feat(ui): link MCP Servers page to the user's connected MCP servers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp-servers/_components/mcp_servers.test.tsx | 14 ++++++++++++++ .../mcp-servers/_components/mcp_servers.tsx | 11 +++++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 1394a923174..091b2f1403f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -62,6 +62,20 @@ describe("MCPServers", () => { expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); + it.each(["Admin", "Internal User"])("links a %s to their MCP connections page", async (userRole) => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); + + render( + + + , + ); + + const myConnections = await screen.findByRole("link", { name: "My Connections" }); + expect(myConnections).toBeVisible(); + expect(myConnections).toHaveAttribute("href", "/ui/connect"); + }); + it("should render mocked MCP servers data in the table", async () => { // Mock MCP servers data const mockServers = [ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index aa0a031c55c..33b2f344bba 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -1,7 +1,8 @@ import { isAdminRole, isProxyAdminRole, isProxyAdminTierRole } from "@/utils/roles"; -import { CircleHelp, Search } from "lucide-react"; +import { CircleHelp, Plug, Search } from "lucide-react"; +import Link from "next/link"; import { Badge } from "@/components/ui/badge"; -import { Button } from "@/components/ui/button"; +import { Button, buttonVariants } from "@/components/ui/button"; import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -43,6 +44,8 @@ import MCPDiscovery from "./mcp_discovery"; import { ByokCredentialModal } from "@/components/mcp_tools/ByokCredentialModal"; import { getSecureItem } from "@/utils/secureStorage"; import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils"; +import { uiHref } from "@/utils/uiHref"; +import { cn } from "@/lib/cva.config"; import UserEnvVarsModal from "./UserEnvVarsModal"; import { listMCPUserEnvVarStatus } from "@/components/networking"; @@ -485,6 +488,10 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i

Configure and manage your MCP servers

+ + + My Connections + {isAdminRole(userRole) && ( <>

Configure and manage your MCP servers

-
+
My Connections From 417a88daedf6e5e4d2de62fab5a33f4ec5a6ae3f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:13:07 -0700 Subject: [PATCH 083/102] fix(responses): carry dict-valued reasoning_effort and keep the frame type on websocket defaults A deployment whose reasoning_effort is an object is copied through as reasoning the way the HTTP mapper does it instead of being dropped, and the relay re-asserts the response.create frame type after merging extra_body so a type key inside it can never replace it. The lazy OpenAPI snapshot goes back to main: the earlier regeneration came from a Python 3.14 interpreter dedenting docstrings, which CI on 3.12 rejects --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- litellm/responses/main.py | 22 +++++++----- litellm/responses/streaming_iterator.py | 2 +- .../test_responses_websocket_all_providers.py | 36 +++++++++++++++++++ 4 files changed, 51 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 2aa2cf15ac1..8a8d08c6887 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19394,7 +19394,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 6064fa66d91..a5912bb42b1 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -2262,24 +2262,28 @@ def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: return metadata -_EXTRA_BODY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object] | None) + + +def _deployment_reasoning_default(kwargs: Mapping[str, object]) -> Reasoning | dict[str, object] | None: + if kwargs.get("reasoning") is not None: + return None + reasoning_effort: Final = kwargs.get("reasoning_effort") + if isinstance(reasoning_effort, str): + return LiteLLMResponsesTransformationHandler()._map_reasoning_effort(reasoning_effort) + return _JSON_OBJECT_ADAPTER.validate_python(reasoning_effort) if isinstance(reasoning_effort, Mapping) else None def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults: - reasoning_effort: Final = kwargs.get("reasoning_effort") - mapped_reasoning: Final = ( - LiteLLMResponsesTransformationHandler()._map_reasoning_effort(reasoning_effort) - if kwargs.get("reasoning") is None and isinstance(reasoning_effort, str) - else None - ) + default_reasoning: Final = _deployment_reasoning_default(kwargs) candidate_params: Final[dict[str, object]] = { **kwargs, - **({"reasoning": mapped_reasoning} if mapped_reasoning is not None else {}), + **({"reasoning": default_reasoning} if default_reasoning is not None else {}), } fill_missing: Final = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(candidate_params) return ResponsesWebSocketRequestDefaults( fill_missing=MappingProxyType(dict(fill_missing)), - overrides=MappingProxyType(_EXTRA_BODY_ADAPTER.validate_python(kwargs.get("extra_body")) or {}), + overrides=MappingProxyType(_JSON_OBJECT_ADAPTER.validate_python(kwargs.get("extra_body")) or {}), ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 1c82df8c664..8f65552ec20 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1883,7 +1883,7 @@ class ResponsesWebSocketStreaming: nested: Final = msg_obj.get("response") if _is_json_object(nested): return {**msg_obj, "response": self.request_defaults.merged_into(nested)} - return self.request_defaults.merged_into(msg_obj) + return {**self.request_defaults.merged_into(msg_obj), "type": msg_obj["type"]} async def _mask_response_create(self, message: str) -> str: """ diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 918c4a33d40..b05904acdcb 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1252,6 +1252,42 @@ class TestNativeWebSocketDeploymentDefaults: assert dict(defaults.fill_missing) == {"reasoning": {"effort": "low"}} assert dict(defaults.overrides) == {} + def test_builder_copies_dict_valued_reasoning_effort_like_the_http_path(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + {"model": "gpt-5-pro", "reasoning_effort": {"effort": "xhigh", "summary": "auto"}} + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "xhigh", "summary": "auto"}} + + @pytest.mark.asyncio + async def test_extra_body_type_key_never_replaces_the_frame_type(self): + from types import MappingProxyType + + from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults + + handler = _make_streaming( + authorized_model="gpt-5-pro", + request_defaults=ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType({}), + overrides=MappingProxyType({"type": "session.update", "provider_default": "configured"}), + ), + ) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps({"type": "response.create", "model": "gpt-5-pro", "input": "hi"}) + ) + ) + + assert forwarded == { + "type": "response.create", + "model": "gpt-5-pro", + "input": "hi", + "provider_default": "configured", + } + @pytest.mark.asyncio async def test_flat_frame_gets_defaults_client_keys_win_extra_body_overrides(self): handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) From 1adbfbfbb11b534df9c4c75f5937c51ca2d2ec1b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:28:45 -0700 Subject: [PATCH 084/102] fix: strip eager_input_streaming for non-Claude providers next to input_examples --- litellm/main.py | 46 +++++++++++----------- tests/test_litellm/test_main.py | 70 +++++++++++++++++++++++++-------- 2 files changed, 76 insertions(+), 40 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 859e0df142c..ac8fa507728 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1146,37 +1146,35 @@ def responses_api_bridge_check( return model_info, model -def _should_allow_input_examples(custom_llm_provider: str | None, model: str) -> bool: +_ANTHROPIC_ONLY_TOOL_KEYS: Final = frozenset({"input_examples", "eager_input_streaming"}) + + +def _is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: if custom_llm_provider == "anthropic": return True - if custom_llm_provider == "azure_ai" or custom_llm_provider == "bedrock" or custom_llm_provider == "vertex_ai": - return "claude" in model.lower() + model_lower: Final = model.lower() + if custom_llm_provider == "bedrock": + return "claude" in model_lower or ("arn:" in model_lower and ":bedrock:" in model_lower) + if custom_llm_provider == "azure_ai" or custom_llm_provider == "vertex_ai": + return "claude" in model_lower return False -def _drop_input_examples_from_tool(tool: dict) -> dict: - tool_copy: Final = tool.copy() - tool_copy.pop("input_examples", None) - function = tool_copy.get("function") - if isinstance(function, dict): - function = function.copy() - function.pop("input_examples", None) - tool_copy["function"] = function - return tool_copy +def _without_anthropic_only_tool_keys(tool: dict) -> dict: + kept: Final = {key: value for key, value in tool.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS} + function: Final = tool.get("function") + if not isinstance(function, dict): + return kept + return { + **kept, + "function": {key: value for key, value in function.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS}, + } -def _drop_input_examples_from_tools( - tools: list[dict] | None, -) -> list[dict] | None: +def _drop_anthropic_only_tool_keys(tools: list[dict] | None) -> list[dict] | None: if tools is None: return None - cleaned_tools: Final[list[dict]] = [] - for tool in tools: - if isinstance(tool, dict): - cleaned_tools.append(_drop_input_examples_from_tool(tool)) - else: - cleaned_tools.append(tool) - return cleaned_tools + return [_without_anthropic_only_tool_keys(tool) if isinstance(tool, dict) else tool for tool in tools] class _ProxyAuthHeadersProvider(Protocol): @@ -5360,8 +5358,8 @@ def completion( api_base=api_base, ) - if not _should_allow_input_examples(custom_llm_provider=custom_llm_provider, model=model): - tools = _drop_input_examples_from_tools(tools=tools) + if not _is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): + tools = _drop_anthropic_only_tool_keys(tools=tools) if provider_specific_header is not None: headers.update( diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index cbbac3d247f..bd115c699d5 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -349,28 +349,66 @@ def test_bedrock_latency_optimized_inference(): assert json_data["performanceConfig"]["latency"] == "optimized" -def test_strip_input_examples_for_non_anthropic_providers(): +@pytest.mark.parametrize( + ("custom_llm_provider", "model", "expected"), + [ + ("anthropic", "claude-sonnet-5", True), + ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), + ("bedrock", "us.amazon.nova-2-lite-v1:0", False), + ("vertex_ai", "claude-sonnet-5", True), + ("vertex_ai", "gemini-3.8-flash", False), + ("azure_ai", "claude-sonnet-4-6", True), + ("azure_ai", "gpt-5.6", False), + ("openai", "gpt-5.6", False), + ("gemini", "gemini-3.8-flash", False), + ], +) +def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): + assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected + + +@pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"]) +def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str): tools = [ - { - "type": "function", - "name": "example_tool", - "input_examples": [{"foo": "bar"}], - "function": { - "name": "example_tool", - "input_examples": [{"foo": "bar"}], - }, - } + {"type": "function", "name": "example_tool", key: True, "function": {"name": "example_tool", key: True}}, + "opaque_tool", ] - assert not litellm_main._should_allow_input_examples( - custom_llm_provider="openai", model="gpt-4o-mini" + cleaned = litellm_main._drop_anthropic_only_tool_keys(tools=tools) + + assert cleaned == [ + {"type": "function", "name": "example_tool", "function": {"name": "example_tool"}}, + "opaque_tool", + ] + assert tools[0][key] is True + assert tools[0]["function"][key] is True + + +def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx.MockRouter, openai_api_response): + api_base: Final = "http://localhost:12346/v1" + mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) ) - cleaned = litellm_main._drop_input_examples_from_tools(tools=tools) + litellm.completion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Write the file"}], + tools=[ + { + "type": "function", + "function": {"name": "write_file", "parameters": {"type": "object", "properties": {}}}, + "eager_input_streaming": True, + } + ], + api_base=api_base, + api_key="fake_openai_api_key", + ) - assert isinstance(cleaned, list) - assert "input_examples" not in cleaned[0] - assert "input_examples" not in cleaned[0]["function"] + assert mock_route.called + sent_tool: Final = json.loads(respx_mock.calls[0].request.content)["tools"][0] + assert "eager_input_streaming" not in sent_tool + assert sent_tool["function"]["name"] == "write_file" def test_custom_provider_with_extra_headers(): From 66c01cf35c14d06877f7560af7604a5c8e36154c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:39:02 -0700 Subject: [PATCH 085/102] refactor(responses): map finish reasons to incomplete_details through a lookup table The match statement in _incomplete_details_for_finish_reason tripped CodeQL's mixed explicit and implicit returns alert (code-scanning 12640). A module-level MappingProxyType keyed by finish reason gives the same three mappings with one explicit return path --- .../transformation.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 961fffd3d42..7337beba45c 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -112,6 +112,9 @@ ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) +_INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_tokens", "content_filter"]]] = ( + MappingProxyType({"length": "max_output_tokens", "content_filter": "content_filter", "refusal": "content_filter"}) +) @dataclass(frozen=True, slots=True) @@ -2303,13 +2306,10 @@ class LiteLLMCompletionResponsesConfig: ) -> IncompleteDetails | None: if existing is not None: return existing - match finish_reason: - case "length": - return IncompleteDetails(reason="max_output_tokens") - case "content_filter" | "refusal": - return IncompleteDetails(reason="content_filter") - case _: - return None + if finish_reason is None: + return None + reason: Final = _INCOMPLETE_REASON_BY_FINISH_REASON.get(finish_reason) + return IncompleteDetails(reason=reason) if reason is not None else None @staticmethod def _tool_call_id_from_responses_item(item_id: str | None, call_id: str | None) -> str: From 44034c1d5e25c8f1888dfd3f13e847cc44860fab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:43:07 -0700 Subject: [PATCH 086/102] test: source the gemma context window limits and isolate the cost map cache --- .../gemma/test_vertex_ai_gemma_global_endpoint.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py index 9d08daa0a81..85a2124ab02 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py @@ -180,12 +180,11 @@ class TestCreateVertexURLGemma: # --------------------------------------------------------------------------- -def test_gemma_maas_context_window_matches_google(monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - +def test_gemma_maas_context_window_matches_google(local_model_cost_map): info = litellm.get_model_info("vertex_ai/google/gemma-4-26b-a4b-it-maas") + # 262,144 context length and 128,000 maximum output per Google's model page, checked 2026-09-18: + # https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/maas/google/gemma-4-26b-a4b-it assert info["max_input_tokens"] == 262144 assert info["max_output_tokens"] == 128000 From b4bfd92a2a4ab11df53d886704b621b6e8cd2339 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 13:36:06 -0700 Subject: [PATCH 087/102] refactor(rust): route-neutral callback contract Every legacy callback call from callbacks-legacy now goes through one typed Python shim, litellm.rust_bridge.legacy_callbacks, the only Python module the crate reaches. Before, the crate called Logging methods, litellm.utils hooks, the logging worker, the executor and several litellm globals directly, and its tests retyped those signatures by hand, so an outdated fake could accept a call the real code rejects. python_contract.json lists each shim function's parameters: a Python test pins it to the real signatures and a Rust test pins it to the Rust enum. The lifecycle contract changes to match the Python @client wrapper: - the driver emits CallEvent::Started before begin, so every host sees one start time - RequestContext carries the route-resolved api_key, so legacy pre_call and post_call receive it, and post_call's additional_args match the Python OCR path - Passthrough and its re-aliasing are gone - async deployment hooks always run, and the "no callbacks" shortcut that skipped the logging payload is removed, as in the Python path The OCR api_key is a SecretValue from the wire request onward, so Debug output upstream of the callback contract cannot leak it. host-python's RouteHost now classifies native failures once through classify, and host ops return HostOpError. The OCR route host keeps main's public errors by sending both through the existing Python map_failure. --- litellm-rust/Cargo.lock | 3 + litellm-rust/crates/auth/src/secret.rs | 4 +- .../crates/callbacks-legacy/AGENTS.md | 14 +- .../crates/callbacks-legacy/Cargo.toml | 4 + .../callbacks-legacy/python_contract.json | 98 +++++ .../crates/callbacks-legacy/src/adapter.rs | 101 +++-- .../crates/callbacks-legacy/src/call.rs | 43 +-- .../crates/callbacks-legacy/src/callbacks.rs | 242 +++--------- .../callbacks-legacy/src/legacy_python.rs | 155 ++++++++ .../crates/callbacks-legacy/src/lib.rs | 5 +- .../crates/callbacks-legacy/src/logger.rs | 91 +---- .../callbacks-legacy/src/preparation.rs | 44 +-- .../crates/callbacks-legacy/tests/deferred.rs | 18 +- .../tests/deployment_hooks.rs | 18 +- .../crates/callbacks-legacy/tests/payload.rs | 113 +++--- .../crates/callbacks-legacy/tests/support.rs | 135 ++++--- .../crates/callbacks-legacy/tests/terminal.rs | 46 ++- litellm-rust/crates/callbacks/Cargo.toml | 1 + litellm-rust/crates/callbacks/src/event.rs | 77 +--- litellm-rust/crates/callbacks/src/run.rs | 36 +- litellm-rust/crates/core/src/ocr/handler.rs | 13 +- litellm-rust/crates/core/src/ocr/mod.rs | 4 +- litellm-rust/crates/core/src/ocr/prepare.rs | 4 +- .../crates/core/src/ocr/provider_config.rs | 46 ++- litellm-rust/crates/core/src/ocr/types.rs | 18 +- litellm-rust/crates/core/src/ocr/wire.rs | 4 +- .../tests/azure_document_intelligence_ocr.rs | 2 +- litellm-rust/crates/core/tests/ocr.rs | 22 +- .../crates/core/tests/ocr/document.rs | 152 ++++++++ .../crates/core/tests/ocr/passthrough.rs | 282 -------------- litellm-rust/crates/core/tests/ocr/support.rs | 10 +- litellm-rust/crates/host-python/AGENTS.md | 5 +- .../crates/host-python/src/adapter.rs | 55 ++- .../crates/host-python/src/argument.rs | 51 +++ litellm-rust/crates/host-python/src/driver.rs | 364 +++++++++++++----- litellm-rust/crates/host-python/src/lib.rs | 8 +- .../document_intelligence/transformation.rs | 21 +- .../llms/src/azure_ai/ocr/transformation.rs | 21 +- .../llms/src/base_llm/ocr/transformation.rs | 17 +- .../llms/src/cohere/ocr/transformation.rs | 6 +- .../llms/src/custom_httpx/llm_http_handler.rs | 32 +- .../llms/src/mistral/ocr/transformation.rs | 6 +- .../llms/src/reducto/ocr/transformation.rs | 8 +- .../llms/src/vertex_ai/ocr/transformation.rs | 5 +- litellm-rust/crates/python-bridge/AGENTS.md | 2 +- .../python-bridge/src/routes/ocr/host.rs | 61 +-- .../python-bridge/src/routes/ocr/project.rs | 10 +- litellm/litellm_core_utils/litellm_logging.py | 43 +-- litellm/rust_bridge/legacy_callbacks.py | 291 ++++++++++---- .../rust_bridge/test_legacy_callbacks.py | 19 +- tests/test_litellm_rust/ocr/test_lifecycle.py | 31 +- 51 files changed, 1564 insertions(+), 1297 deletions(-) create mode 100644 litellm-rust/crates/callbacks-legacy/python_contract.json create mode 100644 litellm-rust/crates/callbacks-legacy/src/legacy_python.rs create mode 100644 litellm-rust/crates/core/tests/ocr/document.rs delete mode 100644 litellm-rust/crates/core/tests/ocr/passthrough.rs create mode 100644 litellm-rust/crates/host-python/src/argument.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c359ca19986..0265e0adbc2 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2031,6 +2031,7 @@ dependencies = [ name = "litellm-callbacks" version = "0.1.0" dependencies = [ + "litellm-auth", "rstest", "serde_json", "tokio", @@ -2040,11 +2041,13 @@ dependencies = [ name = "litellm-callbacks-legacy" version = "0.1.0" dependencies = [ + "litellm-auth", "litellm-callbacks", "litellm-host-python", "pyo3", "rstest", "serde_json", + "strum", ] [[package]] diff --git a/litellm-rust/crates/auth/src/secret.rs b/litellm-rust/crates/auth/src/secret.rs index 3ecb0a835ee..a07fe3eaad9 100644 --- a/litellm-rust/crates/auth/src/secret.rs +++ b/litellm-rust/crates/auth/src/secret.rs @@ -1,6 +1,8 @@ +use serde::Deserialize; use veil::Redact; -#[derive(Redact, Clone)] +#[derive(Redact, Clone, Deserialize)] +#[serde(transparent)] pub struct SecretValue(#[redact(with = "[REDACTED]")] String); impl SecretValue { diff --git a/litellm-rust/crates/callbacks-legacy/AGENTS.md b/litellm-rust/crates/callbacks-legacy/AGENTS.md index e4762d3037a..e184e2fb415 100644 --- a/litellm-rust/crates/callbacks-legacy/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy/AGENTS.md @@ -1,15 +1,17 @@ - Target invariants, not completion claims - Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits) - - The driver in `litellm-host-python`, the routes and core see one `CallbackAdapter`; they never learn which Python objects consume a call + - The driver in `litellm-host-python`, the routes and core see one `PythonLifecycle`; they never learn which Python objects consume a call +- Rust drives the call; every litellm Python internal it still borrows is a variant of `LegacyPython`, grouped by subsystem (`Wrapper`, `Logging`, `DeploymentHooks`) + - The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it + - Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython` - `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy -- `setup` decides once who owns the `Logging` instance and returns it as `CallSetup.bridge_owned`; `PythonLogger` carries it and nothing on the instance records it - - A logger the caller passed as `litellm_logging_obj` is caller-owned and observed in full, because the caller reads it after the call; the proxy is the live case - - A logger `function_setup` built for this call is bridge-owned, so each fan-out phase is skipped when `callbacks_needed` finds no registry, dynamic callback, `logger_fn` or debug switch for it; cost, timing and response metadata still run +- `setup` reuses a `Logging` the caller passed as `litellm_logging_obj` (the proxy and Router are the live cases) and otherwise builds one through `function_setup`, as `@client` does + - Either way every phase calls the same `Logging` method the Python path calls; which callbacks run is `Logging`'s decision, never this crate's - Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view - - Re-alias every `passthrough_fields` body key to the caller's object before `pre_call`; a keyword wins over the request attribute even when it is an explicit `None` + - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup` - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only - - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only facts that cross from the route are the prepared keyword view and `RequestContext.passthrough_fields` + - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view - Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy/Cargo.toml index 96c9c9ed560..3cf9382f857 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy/Cargo.toml @@ -9,8 +9,12 @@ autotests = false [dependencies] litellm-callbacks.workspace = true litellm-host-python.workspace = true + pyo3.workspace = true +strum.workspace = true +serde_json.workspace = true [dev-dependencies] +litellm-auth.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy/python_contract.json new file mode 100644 index 00000000000..840c0abfa45 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy/python_contract.json @@ -0,0 +1,98 @@ +{ + "setup": [ + "call_type", + "args", + "kwargs", + "start_time", + "asynchronous" + ], + "check_limits": [ + "kwargs" + ], + "finalize": [ + "response", + "logger", + "kwargs", + "start_time", + "end_time" + ], + "update_logging": [ + "logger", + "kwargs", + "model", + "optional_params", + "litellm_params", + "custom_llm_provider" + ], + "pre_call": [ + "logger", + "input", + "api_key", + "additional_args" + ], + "post_call": [ + "logger", + "original_response", + "api_key", + "additional_args" + ], + "defers_async_logging": [ + "logger" + ], + "defer_success": [ + "logger", + "pending" + ], + "sync_success_for_async_call": [ + "logger", + "response", + "start", + "end" + ], + "failure_handler": [ + "logger", + "error", + "start", + "end", + "asynchronous" + ], + "submit_success": [ + "logger", + "response", + "start", + "end" + ], + "async_success_handler": [ + "logger", + "response", + "start", + "end" + ], + "enqueue_logging": [ + "coroutine" + ], + "restore_context": [ + "logger" + ], + "custom_pricing_fields": [], + "is_internal_call": [], + "credential_list": [], + "warn_unknown_credential": [ + "name", + "loaded" + ], + "before_deployment_call": [ + "kwargs", + "call_type" + ], + "after_deployment_success": [ + "kwargs", + "response", + "call_type" + ], + "after_deployment_failure": [ + "kwargs", + "error", + "call_type" + ] +} diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index df346506094..7204207cd62 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -4,7 +4,7 @@ use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest}; use litellm_host_python::{ - AdapterStep, CallbackAdapter, PublicValue, from_py, missing_state, to_py, + LifecycleStep, PublicValue, PythonLifecycle, from_py, missing_state, to_py, }; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -12,6 +12,7 @@ use pyo3::{ prelude::*, types::PyDict, }; +use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, @@ -43,7 +44,7 @@ pub struct LegacyLogging { response: Option>, error: Option>, body: Option>, - headers: Option>, + context: Option, asynchronous: bool, internal: bool, pending: Option, @@ -76,7 +77,7 @@ impl LegacyLogging { response: None, error: None, body: None, - headers: None, + context: None, asynchronous, internal: false, pending: None, @@ -85,8 +86,8 @@ impl LegacyLogging { /// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never /// runs them. - fn deployment_hooks(&self, py: Python<'_>) -> PyResult { - Ok(self.asynchronous && DeploymentHooks::needed(py)?) + fn runs_deployment_hooks(&self) -> bool { + self.asynchronous } fn logger(&self) -> PyResult<&PythonLogger> { @@ -95,13 +96,13 @@ impl LegacyLogging { }) } - fn prepare(&mut self, py: Python<'_>) -> PyResult { + fn prepare(&mut self, py: Python<'_>) -> PyResult { let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind(); self.call.set_kwargs(prepared); - Ok(AdapterStep::Arguments(self.call.kwargs().clone_ref(py))) + Ok(LifecycleStep::Arguments(self.call.kwargs().clone_ref(py))) } - fn finalize(&mut self, py: Python<'_>) -> PyResult { + fn finalize(&mut self, py: Python<'_>) -> PyResult { finalize( py, &self.response, @@ -112,7 +113,7 @@ impl LegacyLogging { )?; self.response .as_ref() - .map(|response| AdapterStep::Response(response.clone_ref(py))) + .map(|response| LifecycleStep::Response(response.clone_ref(py))) .ok_or_else(missing_state) } @@ -145,9 +146,7 @@ impl LegacyLogging { .get_item("fallbacks")? .is_none_or(|value| value.is_none()) { - if !logger.callbacks_needed(py, "async_success")? { - logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?; - } else if logger.defers_async_logging(py) { + if logger.defers_async_logging(py) { let pending = Py::new( py, PendingLogging { @@ -165,12 +164,12 @@ impl LegacyLogging { /// The sync failure handler, then the async one for async calls. Ordinary handler /// errors never replace the selected failure or suppress the other family; a /// cancellation does end the call. - fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult { + fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult { let (Some(logger), Some(error)) = (&self.logger, &self.error) else { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); }; if self.asynchronous && self.internal { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); } if let Err(failure) = logger.failure(py, error, &self.start, &self.end, false) && is_cancellation(py, &failure) @@ -178,27 +177,27 @@ impl LegacyLogging { return Err(failure); } if !self.asynchronous { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); } match logger.failure(py, error, &self.start, &self.end, true) { Ok(Some(awaitable)) => { self.pending = Some(Pending::AsyncFailure); - Ok(AdapterStep::Await(awaitable)) + Ok(LifecycleStep::Await(awaitable)) } - Ok(None) => Ok(AdapterStep::Done), + Ok(None) => Ok(LifecycleStep::Done), Err(failure) if is_cancellation(py, &failure) => Err(failure), - Err(_) => Ok(AdapterStep::Done), + Err(_) => Ok(LifecycleStep::Done), } } } -impl CallbackAdapter for LegacyLogging { +impl PythonLifecycle for LegacyLogging { fn begin( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult { + ) -> PyResult { self.call.set_kwargs(arguments); self.start = datetime(py, started_at)?; self.internal = is_internal_call(py)?; @@ -212,9 +211,9 @@ impl CallbackAdapter for LegacyLogging { )?; self.logger = Some(result.logger()?); self.call.set_kwargs(result.kwargs()?); - if self.deployment_hooks(py)? { + if self.runs_deployment_hooks() { self.pending = Some(Pending::DeploymentPreCall); - return Ok(AdapterStep::Await(DeploymentHooks::before_call( + return Ok(LifecycleStep::Await(DeploymentHooks::before_call( py, self.call.kwargs(), self.surface.call_type, @@ -228,18 +227,16 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult { + ) -> PyResult { let logger = self.logger()?; logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?; - if !logger.callbacks_needed(py, "payload")? { - logger.record_api_call_start(py)?; - return Ok(AdapterStep::Wire(wire)); - } let body = to_py(py, &wire.body)? .into_bound(py) .cast_into::()?; - for name in context.passthrough_fields.iter() { - if let Some(value) = self.call.lookup(py, name)? { + for (name, sent) in wire.body.as_object().into_iter().flatten() { + if let Some(value) = self.call.lookup(py, name)? + && from_py::(&value).is_ok_and(|caller| caller == *sent) + { body.set_item(name, value)?; } } @@ -248,12 +245,11 @@ impl CallbackAdapter for LegacyLogging { headers.set_item(name, value)?; } self.body = Some(body.clone().unbind()); - self.headers = Some(headers.clone().unbind()); - let api_key = self.call.lookup(py, "api_key")?; + self.context = Some(context.clone()); self.logger()?.pre_call( py, self.surface.input_description, - api_key.as_ref(), + context.api_key.as_ref().map(|api_key| api_key.expose()), &body, &headers, &wire.url, @@ -262,7 +258,7 @@ impl CallbackAdapter for LegacyLogging { .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) .collect::>>()?; - Ok(AdapterStep::Wire(Box::new(WireRequest { + Ok(LifecycleStep::Wire(Box::new(WireRequest { body: from_py(&body)?, headers, ..*wire @@ -274,12 +270,12 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult { + ) -> PyResult { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response); - if self.deployment_hooks(py)? { + if self.runs_deployment_hooks() { self.pending = Some(Pending::DeploymentPostCall); - return Ok(AdapterStep::Await(DeploymentHooks::after_success( + return Ok(LifecycleStep::Await(DeploymentHooks::after_success( py, self.call.kwargs(), &self.response, @@ -294,31 +290,35 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult { + ) -> PyResult { match (event, public) { + (CallEvent::Started { .. }, _) => Ok(LifecycleStep::Done), (CallEvent::ResponseReceived { raw }, _) => { - let logger = self.logger()?; - if logger.callbacks_needed(py, "payload")? { - logger.post_call(py, &raw.body, self.body.as_ref(), self.headers.as_ref())?; - } - Ok(AdapterStep::Done) + let api_key = self + .context + .as_ref() + .and_then(|context| context.api_key.as_ref()) + .map(|api_key| api_key.expose()); + self.logger()? + .post_call(py, &raw.body, api_key, self.body.as_ref())?; + Ok(LifecycleStep::Done) } (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); self.dispatch_success(py)?; - Ok(AdapterStep::Done) + Ok(LifecycleStep::Done) } (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); if *origin == FailureOrigin::Call && self.logger.is_some() - && self.deployment_hooks(py)? + && self.runs_deployment_hooks() { let error = self.error.as_ref().ok_or_else(missing_state)?; self.pending = Some(Pending::DeploymentFailure); - return Ok(AdapterStep::Await(DeploymentHooks::after_failure( + return Ok(LifecycleStep::Await(DeploymentHooks::after_failure( py, self.call.kwargs(), error, @@ -331,7 +331,7 @@ impl CallbackAdapter for LegacyLogging { } } - fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { match self.pending.take().ok_or_else(missing_state)? { Pending::DeploymentPreCall => { self.call @@ -345,7 +345,7 @@ impl CallbackAdapter for LegacyLogging { Pending::DeploymentFailure => self.dispatch_failure(py), Pending::AsyncFailure => match result { Err(failure) if is_cancellation(py, &failure) => Err(failure), - _ => Ok(AdapterStep::Done), + _ => Ok(LifecycleStep::Done), }, } } @@ -357,7 +357,7 @@ impl CallbackAdapter for LegacyLogging { error.write_unraisable(py, None); } self.body = None; - self.headers = None; + self.context = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -369,8 +369,7 @@ impl CallbackAdapter for LegacyLogging { visit.call(&self.end)?; visit.call(&self.response)?; visit.call(&self.error)?; - visit.call(&self.body)?; - visit.call(&self.headers) + visit.call(&self.body) } } diff --git a/litellm-rust/crates/callbacks-legacy/src/call.rs b/litellm-rust/crates/callbacks-legacy/src/call.rs index 59090ee8d60..bd1e2525d3d 100644 --- a/litellm-rust/crates/callbacks-legacy/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy/src/call.rs @@ -4,7 +4,7 @@ //! this crate holds them. use litellm_callbacks::{machine::Machine, route::Route}; -use litellm_host_python::{RouteHost, run_call}; +use litellm_host_python::{RouteHost, lookup, run_call}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -63,21 +63,6 @@ impl PublicCall { } } -/// The caller's own object for a public argument, as every legacy reader resolves it: the -/// keyword if given, even an explicit `None`, else the bound request's attribute. A route -/// host projecting from the prepared keyword view uses the same rule, so the callbacks -/// and the provider see one object per argument. -pub fn lookup<'py>( - kwargs: &Bound<'py, PyDict>, - request: &Bound<'py, PyAny>, - name: &str, -) -> PyResult>> { - if let Some(value) = kwargs.get_item(name)? { - return Ok(Some(value)); - } - request.getattr_opt(name) -} - /// Runs one native call under the legacy `Logging` contract: the route host projects from /// the keyword view the contract prepares, and the contract observes the call. pub fn run_legacy_call( @@ -121,32 +106,6 @@ mod tests { (call, locals) } - #[test] - fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { - Python::initialize(); - Python::attach(|py| { - let (call, locals) = capture( - py, - c" -key = object() -document = {'type': 'document_url'} -class Request: - api_key = 'from-request' - api_base = 'from-request' - document = document -request = Request() -kwargs = {'api_key': key, 'api_base': None} -", - ); - let key = locals.get_item("key").unwrap().unwrap(); - let document = locals.get_item("document").unwrap().unwrap(); - assert!(call.lookup(py, "api_key").unwrap().unwrap().is(&key)); - assert!(call.lookup(py, "api_base").unwrap().unwrap().is_none()); - assert!(call.lookup(py, "document").unwrap().unwrap().is(&document)); - assert!(call.lookup(py, "model").unwrap().is_none()); - }); - } - #[test] fn capture_copies_the_keyword_dict_without_copying_its_values() { Python::initialize(); diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs index aa586013e75..9464f1d6612 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs @@ -6,11 +6,10 @@ use litellm_callbacks::event::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; +use crate::legacy_python::{Logging, Wrapper}; use crate::logger::PythonLogger; pub trait LegacyCallbacks { - fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult; - /// `Logging.update_from_kwargs`: what the logger is told about the request it is /// about to see, with consumed credentials redacted. fn update_from_kwargs( @@ -21,26 +20,24 @@ pub trait LegacyCallbacks { context: &RequestContext, ) -> PyResult<()>; - fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()>; - - /// `Logging.pre_call`, or its payload-free shortcut when no input callback listens. + /// `Logging.pre_call`. fn pre_call( &self, py: Python<'_>, input: &str, - api_key: Option<&Bound<'_, PyAny>>, + api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, ) -> PyResult<()>; - /// `Logging.post_call`, or its payload-free shortcut when no input callback listens. + /// `Logging.post_call`. fn post_call( &self, py: Python<'_>, original_response: &str, + api_key: Option<&str>, body: Option<&Py>, - headers: Option<&Py>, ) -> PyResult<()>; fn defers_async_logging(&self, py: Python<'_>) -> bool; @@ -82,16 +79,6 @@ pub trait LegacyCallbacks { } impl LegacyCallbacks for PythonLogger { - fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult { - if !self.bridge_owned() { - return Ok(true); - } - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("callbacks_needed")? - .call1((self.object(py), phase))? - .extract() - } - fn update_from_kwargs( &self, py: Python<'_>, @@ -100,18 +87,13 @@ impl LegacyCallbacks for PythonLogger { context: &RequestContext, ) -> PyResult<()> { let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect(); - let update = PyDict::new(py); - update.set_item("kwargs", redact(py, kwargs.bind(py), &secret_fields)?)?; - update.set_item("model", &context.model)?; - update.set_item( - "optional_params", - redact( - py, - &to_py(py, &context.optional_params)? - .into_bound(py) - .cast_into::()?, - &secret_fields, - )?, + let redacted_kwargs = redact(py, kwargs.bind(py), &secret_fields)?; + let optional_params = redact( + py, + &to_py(py, &context.optional_params)? + .into_bound(py) + .cast_into::()?, + &secret_fields, )?; let params = PyDict::new(py); params.set_item( @@ -131,15 +113,17 @@ impl LegacyCallbacks for PythonLogger { params.set_item(name, value)?; } } - update.set_item("litellm_params", params)?; - update.set_item("custom_llm_provider", &context.custom_llm_provider)?; - self.object(py) - .call_method("update_from_kwargs", (), Some(&update))?; - Ok(()) - } - - fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()> { - self.object(py).call_method0("record_api_call_start_time")?; + Logging::Update.call( + py, + ( + self.object(py), + redacted_kwargs, + &context.model, + optional_params, + params, + &context.custom_llm_provider, + ), + )?; Ok(()) } @@ -147,7 +131,7 @@ impl LegacyCallbacks for PythonLogger { &self, py: Python<'_>, input: &str, - api_key: Option<&Bound<'_, PyAny>>, + api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, @@ -156,17 +140,7 @@ impl LegacyCallbacks for PythonLogger { additional.set_item("complete_input_dict", body)?; additional.set_item("headers", headers)?; additional.set_item("api_base", url)?; - let kwargs = PyDict::new(py); - kwargs.set_item("input", input)?; - kwargs.set_item("api_key", api_key)?; - kwargs.set_item("additional_args", &additional)?; - if self.callbacks_needed(py, "input")? { - self.object(py).call_method("pre_call", (), Some(&kwargs))?; - } else { - self.object(py) - .call_method("_pre_call", (), Some(&kwargs))?; - self.record_api_call_start(py)?; - } + Logging::PreCall.call(py, (self.object(py), input, api_key, &additional))?; Ok(()) } @@ -174,37 +148,28 @@ impl LegacyCallbacks for PythonLogger { &self, py: Python<'_>, original_response: &str, + api_key: Option<&str>, body: Option<&Py>, - headers: Option<&Py>, ) -> PyResult<()> { let additional = PyDict::new(py); additional.set_item("complete_input_dict", body)?; - additional.set_item("headers", headers)?; - if self.callbacks_needed(py, "input")? { - let kwargs = PyDict::new(py); - kwargs.set_item("original_response", original_response)?; - kwargs.set_item("additional_args", &additional)?; - self.object(py) - .call_method("post_call", (), Some(&kwargs))?; - } else { - let response = py - .import("json")? - .call_method1("dumps", (original_response,))?; - self.object(py).call_method1( - "record_post_call", - (response, py.None(), py.None(), additional), - )?; - } + Logging::PostCall.call( + py, + (self.object(py), original_response, api_key, &additional), + )?; Ok(()) } + fn defers_async_logging(&self, py: Python<'_>) -> bool { - self.object(py) - .getattr("_defer_async_logging") - .is_ok_and(|value| value.is_truthy().unwrap_or(false)) + Logging::DefersAsync + .call(py, (self.object(py),)) + .and_then(|value| value.extract()) + .unwrap_or(false) } fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> { - self.object(py).setattr("_native_pending_logging", pending) + Logging::DeferSuccess.call(py, (self.object(py), pending))?; + Ok(()) } fn sync_success_for_async_call( @@ -214,13 +179,7 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success_async")? { - return Ok(()); - } - self.object(py).call_method1( - "handle_sync_success_callbacks_for_async_calls", - (response, start, end), - )?; + Logging::SyncSuccessForAsyncCall.call(py, (self.object(py), response, start, end))?; Ok(()) } @@ -232,34 +191,11 @@ impl LegacyCallbacks for PythonLogger { end: &Option>, asynchronous: bool, ) -> PyResult>> { - if !self.callbacks_needed( - py, - if asynchronous { - "async_failure" - } else { - "sync_failure" - }, - )? { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("failure_bookkeeping")? - .call1((self.object(py), error, start, end, asynchronous))?; - return Ok(None); - } - let trace = py - .import("traceback")? - .getattr("format_exception")? - .call1((error,))?; - let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?; - let value = self.object(py).call_method1( - if asynchronous { - "async_failure_handler" - } else { - "failure_handler" - }, - (error, trace, start, end), - )?; + let value = + Logging::FailureHandler.call(py, (self.object(py), error, start, end, asynchronous))?; Ok(asynchronous.then(|| value.unbind())) } + fn submit_success( &self, py: Python<'_>, @@ -267,22 +203,7 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success")? { - return self.success_bookkeeping(py, response, start, end, false); - } - let context = py.import("contextvars")?.call_method0("copy_context")?; - py.import("litellm.litellm_core_utils.litellm_logging")? - .getattr("executor")? - .call_method1( - "submit", - ( - context.getattr("run")?, - self.object(py).getattr("success_handler")?, - response, - start, - end, - ), - )?; + Logging::SubmitSuccess.call(py, (self.object(py), response, start, end))?; Ok(()) } @@ -293,18 +214,9 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "async_success")? { - return self.success_bookkeeping(py, response, start, end, true); - } - let context = py.import("contextvars")?.call_method0("copy_context")?; - let worker = py - .import("litellm.litellm_core_utils.logging_worker")? - .getattr("GLOBAL_LOGGING_WORKER")? - .getattr("ensure_initialized_and_enqueue")?; - let coroutine = self - .object(py) - .call_method1("async_success_handler", (response, start, end))?; - let enqueue = context.call_method1("run", (worker, &coroutine)); + let coroutine = + Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?; + let enqueue = Logging::Enqueue.call(py, (&coroutine,)); if enqueue.is_err() && let Err(error) = coroutine.call_method0("close") { @@ -315,14 +227,7 @@ impl LegacyCallbacks for PythonLogger { } fn custom_pricing_fields(py: Python<'_>) -> PyResult> { - py.import("litellm.types.utils")? - .getattr("CustomPricingLiteLLMParams")? - .getattr("model_fields")? - .cast_into::()? - .keys() - .iter() - .map(|name| name.extract::()) - .collect() + Logging::CustomPricingFields.call(py, ())?.extract() } fn redact( @@ -347,58 +252,5 @@ fn redact( /// Proxy-internal calls skip the legacy success fan-out. pub fn is_internal_call(py: Python<'_>) -> PyResult { - py.import("litellm._internal_context")? - .getattr("is_internal_call")? - .call_method0("get")? - .extract() -} - -#[cfg(test)] -mod tests { - use pyo3::types::PyDict; - - use super::*; - - fn logger_whose_registries_need_no_input(py: Python<'_>, bridge_owned: bool) -> PythonLogger { - let locals = PyDict::new(py); - py.run( - c" -import sys -import types -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): - sys.modules.setdefault(name, types.ModuleType(name)) -legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] -legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True) -class Logger: - needed = {'input': False} -logger = Logger() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - PythonLogger::new( - locals.get_item("logger").unwrap().unwrap().unbind(), - bridge_owned, - ) - } - - #[test] - fn a_caller_owned_logger_is_observed_in_full() { - Python::initialize(); - Python::attach(|py| { - let logger = logger_whose_registries_need_no_input(py, false); - assert!(logger.callbacks_needed(py, "input").unwrap()); - }); - } - - #[test] - fn a_bridge_owned_logger_is_elided_where_no_registry_needs_it() { - Python::initialize(); - Python::attach(|py| { - let logger = logger_whose_registries_need_no_input(py, true); - assert!(!logger.callbacks_needed(py, "input").unwrap()); - assert!(logger.callbacks_needed(py, "payload").unwrap()); - }); - } + Wrapper::IsInternalCall.call(py, ())?.extract() } diff --git a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs new file mode 100644 index 00000000000..a924775070c --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs @@ -0,0 +1,155 @@ +use pyo3::prelude::*; +use strum::{IntoStaticStr, VariantArray}; + +const MODULE: &str = "litellm.rust_bridge.legacy_callbacks"; + +/// Every litellm Python internal the native call still borrows, grouped by the subsystem it +/// belongs to. Rust drives the call; these exist only so behaviour that Python owns today +/// (span tracking, the standard logging payload, spend, callback fan-out) keeps working. +/// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a +/// user's own callback is not borrowing and does not belong here. +/// +/// `litellm/rust_bridge/legacy_callbacks.py` is the only Python module behind it, and +/// `python_contract.json` pins each function's parameters on both sides. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LegacyPython { + Wrapper(Wrapper), + Logging(Logging), + DeploymentHooks(DeploymentHooks), +} + +/// The `@client` wrapper around the call: `function_setup`, limits, credentials, +/// response metadata and the correlation context. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Wrapper { + #[strum(serialize = "setup")] + Setup, + #[strum(serialize = "check_limits")] + CheckLimits, + #[strum(serialize = "credential_list")] + CredentialList, + #[strum(serialize = "warn_unknown_credential")] + WarnUnknownCredential, + #[strum(serialize = "is_internal_call")] + IsInternalCall, + #[strum(serialize = "finalize")] + Finalize, + #[strum(serialize = "restore_context")] + RestoreContext, +} + +/// litellm's `Logging` object and the sync and async callback fan-out behind it. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Logging { + #[strum(serialize = "custom_pricing_fields")] + CustomPricingFields, + #[strum(serialize = "update_logging")] + Update, + #[strum(serialize = "pre_call")] + PreCall, + #[strum(serialize = "post_call")] + PostCall, + #[strum(serialize = "defers_async_logging")] + DefersAsync, + #[strum(serialize = "defer_success")] + DeferSuccess, + #[strum(serialize = "sync_success_for_async_call")] + SyncSuccessForAsyncCall, + #[strum(serialize = "submit_success")] + SubmitSuccess, + #[strum(serialize = "async_success_handler")] + AsyncSuccessHandler, + #[strum(serialize = "enqueue_logging")] + Enqueue, + #[strum(serialize = "failure_handler")] + FailureHandler, +} + +/// The `litellm.utils` fan-outs that run every callback's deployment hook. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum DeploymentHooks { + #[strum(serialize = "before_deployment_call")] + BeforeDeploymentCall, + #[strum(serialize = "after_deployment_success")] + AfterDeploymentSuccess, + #[strum(serialize = "after_deployment_failure")] + AfterDeploymentFailure, +} + +impl LegacyPython { + fn name(self) -> &'static str { + match self { + Self::Wrapper(function) => function.into(), + Self::Logging(function) => function.into(), + Self::DeploymentHooks(function) => function.into(), + } + } + + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + py.import(MODULE)?.getattr(self.name())?.call1(args) + } +} + +impl Wrapper { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Wrapper(self).call(py, args) + } +} + +impl Logging { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Logging(self).call(py, args) + } +} + +impl DeploymentHooks { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::DeploymentHooks(self).call(py, args) + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use strum::VariantArray; + + use super::{DeploymentHooks, LegacyPython, Logging, Wrapper}; + use crate::test_support::PYTHON_CONTRACT; + + #[test] + fn every_borrowed_function_is_in_the_python_contract() { + let contract: serde_json::Map = + serde_json::from_str(PYTHON_CONTRACT).unwrap(); + let declared: BTreeSet<&str> = contract.keys().map(String::as_str).collect(); + let called: Vec<&str> = Wrapper::VARIANTS + .iter() + .map(|&function| LegacyPython::Wrapper(function)) + .chain( + Logging::VARIANTS + .iter() + .map(|&function| LegacyPython::Logging(function)), + ) + .chain( + DeploymentHooks::VARIANTS + .iter() + .map(|&function| LegacyPython::DeploymentHooks(function)), + ) + .map(LegacyPython::name) + .collect(); + assert_eq!(called.len(), declared.len(), "a function is borrowed twice"); + assert_eq!(called.into_iter().collect::>(), declared); + } +} diff --git a/litellm-rust/crates/callbacks-legacy/src/lib.rs b/litellm-rust/crates/callbacks-legacy/src/lib.rs index 06783ac255d..42ffd545e2b 100644 --- a/litellm-rust/crates/callbacks-legacy/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy/src/lib.rs @@ -2,7 +2,7 @@ //! sync and async callback registries it fans out to, the deployment hooks, the deferred //! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name //! inheritance, budget and retry-count limits). All of it sits behind one -//! [`CallbackAdapter`](litellm_host_python::CallbackAdapter), so the driver, the routes and +//! [`PythonLifecycle`](litellm_host_python::PythonLifecycle), so the driver, the routes and //! core never learn which Python object is on the other end. //! //! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`] @@ -13,6 +13,7 @@ mod adapter; mod call; mod callbacks; mod deferred; +mod legacy_python; mod logger; mod preparation; #[cfg(test)] @@ -21,7 +22,7 @@ mod test_support; pub(crate) use adapter::LegacyLogging; pub use adapter::LegacySurface; -pub use call::{PublicCall, lookup, run_legacy_call}; +pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; pub(crate) use preparation::prepare; diff --git a/litellm-rust/crates/callbacks-legacy/src/logger.rs b/litellm-rust/crates/callbacks-legacy/src/logger.rs index a0e525000b8..061941f05b9 100644 --- a/litellm-rust/crates/callbacks-legacy/src/logger.rs +++ b/litellm-rust/crates/callbacks-legacy/src/logger.rs @@ -5,34 +5,25 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -/// The `Logging` instance one call fans out through, and who owns it. A logger the caller -/// handed in is observed in full, because the caller reads it after the call; one this -/// crate built through `function_setup` is elided wherever no registry needs it. +use crate::legacy_python::{self, Wrapper}; + +/// The `Logging` instance one call fans out through. pub struct PythonLogger { object: Py, - bridge_owned: bool, } impl PythonLogger { - pub(crate) fn new(object: Py, bridge_owned: bool) -> Self { - Self { - object, - bridge_owned, - } + pub(crate) fn new(object: Py) -> Self { + Self { object } } pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> { self.object.bind(py) } - pub(crate) fn bridge_owned(&self) -> bool { - self.bridge_owned - } - pub fn clone_ref(&self, py: Python<'_>) -> Self { Self { object: self.object.clone_ref(py), - bridge_owned: self.bridge_owned, } } @@ -40,34 +31,17 @@ impl PythonLogger { visit.call(&self.object) } - pub fn success_bookkeeping( - &self, - py: Python<'_>, - response: &Option>, - start: &Py, - end: &Option>, - asynchronous: bool, - ) -> PyResult<()> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("success_bookkeeping")? - .call1((self.object(py), response, start, end, asynchronous))?; - Ok(()) - } - pub fn restore_context(&self, py: Python<'_>) -> PyResult<()> { - py.import("litellm.utils")? - .getattr("_restore_correlation_context_if_supported")? - .call1((self.object(py),))?; + Wrapper::RestoreContext.call(py, (self.object(py),))?; Ok(()) } } -/// A bare Python object was not obtained from `setup`, so it is caller-owned. impl FromPyObject<'_, '_> for PythonLogger { type Error = PyErr; fn extract(object: Borrowed<'_, '_, PyAny>) -> PyResult { - Ok(Self::new(object.to_owned().unbind(), false)) + Ok(Self::new(object.to_owned().unbind())) } } @@ -75,9 +49,7 @@ pub struct SetupResult<'py>(Bound<'py, PyAny>); impl SetupResult<'_> { pub fn logger(&self) -> PyResult { - let object = self.0.getattr("logger")?.unbind(); - let bridge_owned = self.0.getattr("bridge_owned")?.extract()?; - Ok(PythonLogger::new(object, bridge_owned)) + Ok(PythonLogger::new(self.0.getattr("logger")?.unbind())) } pub fn kwargs(&self) -> PyResult> { @@ -93,9 +65,8 @@ pub fn setup<'py>( start: &Py, asynchronous: bool, ) -> PyResult> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("setup")? - .call1((call_type, args, kwargs, start, asynchronous)) + Wrapper::Setup + .call(py, (call_type, args, kwargs, start, asynchronous)) .map(SetupResult) } @@ -107,30 +78,20 @@ pub fn finalize( start: &Py, end: &Option>, ) -> PyResult<()> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("finalize")? - .call1((response, logger.object(py), kwargs, start, end))?; + Wrapper::Finalize.call(py, (response, logger.object(py), kwargs, start, end))?; Ok(()) } pub struct DeploymentHooks; impl DeploymentHooks { - pub fn needed(py: Python<'_>) -> PyResult { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("deployment_callbacks_needed")? - .call0()? - .extract() - } - pub fn before_call( py: Python<'_>, kwargs: &Py, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_pre_call_deployment_hook")? - .call1((kwargs, call_type)) + legacy_python::DeploymentHooks::BeforeDeploymentCall + .call(py, (kwargs, call_type)) .map(Bound::unbind) } @@ -140,9 +101,8 @@ impl DeploymentHooks { response: &Option>, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_post_call_success_deployment_hook")? - .call1((kwargs, response, call_type)) + legacy_python::DeploymentHooks::AfterDeploymentSuccess + .call(py, (kwargs, response, call_type)) .map(Bound::unbind) } @@ -152,9 +112,8 @@ impl DeploymentHooks { error: &Py, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_post_call_failure_deployment_hook")? - .call1((kwargs, error, call_type)) + legacy_python::DeploymentHooks::AfterDeploymentFailure + .call(py, (kwargs, error, call_type)) .map(Bound::unbind) } } @@ -185,10 +144,6 @@ class Setup: reads.append('logger') return logger @property - def bridge_owned(self): - reads.append('bridge_owned') - return True - @property def kwargs(self): reads.append('kwargs') return [] @@ -206,7 +161,6 @@ result = Setup() .object(py) .is(locals.get_item("logger").unwrap().unwrap()) ); - assert!(logger.bridge_owned()); assert!( result .kwargs() @@ -220,17 +174,8 @@ result = Setup() .unwrap() .extract::>() .unwrap(), - ["logger", "bridge_owned", "kwargs"] + ["logger", "kwargs"] ); }); } - - #[test] - fn a_logger_extracted_from_a_bare_object_is_caller_owned() { - Python::initialize(); - Python::attach(|py| { - let logger: PythonLogger = py.None().into_bound(py).extract().unwrap(); - assert!(!logger.bridge_owned()); - }); - } } diff --git a/litellm-rust/crates/callbacks-legacy/src/preparation.rs b/litellm-rust/crates/callbacks-legacy/src/preparation.rs index 981b1702f2e..fa1ff9acd4d 100644 --- a/litellm-rust/crates/callbacks-legacy/src/preparation.rs +++ b/litellm-rust/crates/callbacks-legacy/src/preparation.rs @@ -3,6 +3,8 @@ use pyo3::{ types::{PyDict, PyList}, }; +use crate::legacy_python::Wrapper; + struct CredentialEntry<'py>(Bound<'py, PyAny>); impl<'py> CredentialEntry<'py> { @@ -22,18 +24,19 @@ pub fn prepare<'py>( ) -> PyResult> { let arguments = kwargs.copy()?; arguments.set_item("litellm_logging_obj", logger.object(py))?; - let litellm = py.import("litellm")?; - inherit_credentials(py, &litellm, &arguments)?; - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("check_limits")? - .call1((&arguments,))?; + inherit_credentials(py, &arguments, || { + Ok(Wrapper::CredentialList + .call(py, ())? + .cast_into::()?) + })?; + Wrapper::CheckLimits.call(py, (&arguments,))?; Ok(arguments) } -fn inherit_credentials( - py: Python<'_>, - litellm: &Bound<'_, PyModule>, - arguments: &Bound<'_, PyDict>, +fn inherit_credentials<'py>( + py: Python<'py>, + arguments: &Bound<'py, PyDict>, + credential_list: impl FnOnce() -> PyResult>, ) -> PyResult<()> { let Some(requested) = arguments .get_item("litellm_credential_name")? @@ -45,16 +48,13 @@ fn inherit_credentials( return Ok(()); } let requested: String = requested.extract()?; - let credentials = litellm.getattr("credential_list")?.cast_into::()?; + let credentials = credential_list()?; let names = credentials .iter() .map(|credential| CredentialEntry(credential).name()) .collect::>>()?; let Some(index) = names.iter().position(|name| *name == requested) else { - py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1( - "warning", - ("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()), - )?; + Wrapper::WarnUnknownCredential.call(py, (requested, names.len()))?; return Ok(()); }; let selected = CredentialEntry(credentials.get_item(index)?); @@ -80,19 +80,19 @@ mod tests { } fn inherit(py: Python<'_>, locals: &Bound<'_, PyDict>) -> PyResult<()> { - let litellm = PyModule::new(py, "credential_host")?; - litellm.setattr( - "credential_list", - locals.get_item("credentials").unwrap().unwrap(), - )?; inherit_credentials( py, - &litellm, &locals .get_item("arguments") .unwrap() .unwrap() .cast_into::()?, + || { + Ok(locals + .get_item("credentials")? + .unwrap() + .cast_into::()?) + }, ) } @@ -304,11 +304,11 @@ arguments = {'litellm_credential_name': 'ocr-test'} fn falsy_credential_names_return_before_loading_credentials() { Python::initialize(); Python::attach(|py| { - let litellm = PyModule::new(py, "credential_host").unwrap(); for name in [py.None(), py.eval(c"''", None, None).unwrap().unbind()] { let arguments = PyDict::new(py); arguments.set_item("litellm_credential_name", name).unwrap(); - inherit_credentials(py, &litellm, &arguments).unwrap(); + inherit_credentials(py, &arguments, || panic!("credentials must not be loaded")) + .unwrap(); } }); } diff --git a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs b/litellm-rust/crates/callbacks-legacy/tests/deferred.rs index 3daea8840d8..289ea1b2e7f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deferred.rs @@ -16,7 +16,7 @@ fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { py, PendingLogging { pending: Some(PendingSuccess { - logger: PythonLogger::new(local(&locals, "logger").unbind(), true), + logger: PythonLogger::new(local(&locals, "logger").unbind()), response: Some(local(&locals, "response").unbind()), start: py.None(), end: Some(py.None()), @@ -79,22 +79,6 @@ assert logger.calls == [], logger.calls }); } -#[test] -fn a_release_after_the_async_callbacks_went_away_only_keeps_the_books() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c"logger.needed = {'async_success': False}"); - run( - py, - &locals, - c" -pending.release(True) -assert logger.calls == [('success_bookkeeping', True)], logger.calls -", - ); - }); -} - #[rstest] #[case::ordinary_error(c"RuntimeError('queue full')", false)] #[case::cancellation(c"asyncio.CancelledError()", true)] diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index 3ceda4441a7..ea3510de17e 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue}; +use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -24,7 +24,7 @@ fn begin<'py>( py: Python<'py>, locals: &Bound<'py, PyDict>, asynchronous: bool, -) -> (LegacyLogging, AdapterStep) { +) -> (LegacyLogging, LifecycleStep) { let mut logging = legacy_call(py, locals, asynchronous); let kwargs = local(locals, "kwargs") .cast_into::() @@ -34,15 +34,15 @@ fn begin<'py>( (logging, step) } -fn arguments<'py>(py: Python<'py>, step: AdapterStep) -> Bound<'py, PyDict> { - let AdapterStep::Arguments(arguments) = step else { +fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { + let LifecycleStep::Arguments(arguments) = step else { panic!("expected the prepared arguments"); }; arguments.into_bound(py) } -fn awaits_deployment_hook(step: &AdapterStep) -> bool { - matches!(step, AdapterStep::Await(_)) +fn awaits_deployment_hook(step: &LifecycleStep) -> bool { + matches!(step, LifecycleStep::Await(_)) } #[rstest] @@ -121,7 +121,7 @@ logger.hooks = {'pre': lambda kwargs: kwargs} let step = logging .resume(py, Ok(local(&locals, "replacement").unbind())) .unwrap(); - let AdapterStep::Response(returned) = step else { + let LifecycleStep::Response(returned) = step else { panic!("expected the finalized response"); }; assert!(returned.bind(py).is(local(&locals, "replacement"))); @@ -195,7 +195,7 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle }; assert!(matches!( logging.resume(py, hook_result).unwrap(), - AdapterStep::Await(_) + LifecycleStep::Await(_) )); run( py, @@ -237,7 +237,7 @@ kwargs = {'logger': logger} .unwrap() .unbind(); let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { - AdapterStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), + LifecycleStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), step => Ok(step), }); let error = result.err().unwrap(); diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 480bedf8548..68c0a2b1e15 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,7 +1,8 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{AdapterStep, CallbackAdapter}; +use litellm_auth::SecretValue; +use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host_python::{LifecycleStep, PythonLifecycle}; use pyo3::prelude::*; use rstest::rstest; use serde_json::{Value, json}; @@ -23,20 +24,12 @@ class PayloadLogger(StubLogger): def pre_call(self, input, api_key, additional_args): self.record('pre_call', None) self.pre = additional_args + self.pre_api_key = api_key on_pre_call(additional_args) - def _pre_call(self, input, api_key, additional_args): - self.record('_pre_call', None) - - def record_api_call_start_time(self): - self.record('record_api_call_start_time', None) - - def post_call(self, original_response, additional_args): + def post_call(self, original_response, api_key, additional_args): self.record('post_call', None) - self.post = (original_response, additional_args) - - def record_post_call(self, response, *rest): - self.record('record_post_call', response) + self.post = (original_response, api_key, additional_args) request = Request() kwargs = {} @@ -52,16 +45,16 @@ fn document(source: &str) -> Value { json!({"type": "document_url", "document_url": source}) } -fn before_send(script: &CStr, caller: Value, body: Value) -> WireRequest { - before_send_with_secrets(script, caller, body, &[]) +fn before_send(script: &CStr, body: Value) -> WireRequest { + before_send_with_secrets(script, json!({}), body, &[]) } -/// Runs `before_send` over `body` for a caller whose route-side view is `caller`, with the -/// Python objects `script` binds, then delivers the provider's raw response the way the +/// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with +/// the Python objects `script` binds, then delivers the provider's raw response the way the /// driver does and runs the script's `check()`. fn before_send_with_secrets( script: &CStr, - caller: Value, + optional_params: Value, body: Value, secret_fields: &[&str], ) -> WireRequest { @@ -70,15 +63,15 @@ fn before_send_with_secrets( let locals = namespace(py, PAYLOAD_LOGGER); run(py, &locals, script); let mut logging = LegacyLogging { - logger: Some(PythonLogger::new(local(&locals, "logger").unbind(), true)), + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), ..legacy_call(py, &locals, false) }; let context = RequestContext { model: "model".into(), custom_llm_provider: "provider".into(), - optional_params: caller.clone(), - passthrough_fields: Passthrough::unchanged(caller.as_object().unwrap(), &body), + optional_params, secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), + api_key: Some(SecretValue::new("route-key")), }; let wire = WireRequest { url: "https://provider.invalid/ocr".into(), @@ -93,10 +86,10 @@ fn before_send_with_secrets( }; assert!(matches!( logging.emit(py, &raw, None).unwrap(), - AdapterStep::Done + LifecycleStep::Done )); run(py, &locals, c"check()"); - let AdapterStep::Wire(wire) = step else { + let LifecycleStep::Wire(wire) = step else { panic!("before_send did not hand back the wire request"); }; *wire @@ -129,11 +122,7 @@ def check(): ")] fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); - let wire = before_send( - script, - json!({"document": document(DOCUMENT), "pages": [0]}), - body.clone(), - ); + let wire = before_send(script, body.clone()); assert_eq!(wire.body, body); } @@ -149,7 +138,6 @@ def check(): assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' ", json!({"document": document(DOCUMENT)}), - json!({"document": document(DOCUMENT)}), ); assert_eq!(wire.body["document"], document(EDITED)); } @@ -168,7 +156,6 @@ def check(): assert observed == [False], observed assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} ", - json!({"document": document("https://example.invalid/scan.pdf")}), json!({"document": document(DOCUMENT)}), ); assert_eq!( @@ -177,6 +164,23 @@ def check(): ); } +#[test] +fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { + let body = json!({"pages": [0]}); + let wire = before_send( + c" +opaque = object() +kwargs = {'pages': opaque} +observed = [] +on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) +def check(): + assert observed == [[0]], observed +", + body.clone(), + ); + assert_eq!(wire.body, body); +} + #[rstest] #[case::body( c" @@ -192,7 +196,7 @@ def on_pre_call(args): )] fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, json!({}), body.clone()); + let wire = before_send(script, body.clone()); assert_eq!(wire.body, body); assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); } @@ -205,7 +209,6 @@ def on_pre_call(args): args['headers']['x-callback'] = 'edited' ", json!({}), - json!({}), ); assert_eq!( wire.headers, @@ -288,7 +291,7 @@ def on_pre_call(args): )] fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, json!({"document": document(DOCUMENT)}), body); + let wire = before_send(script, body); assert_eq!(wire.body, expected); } @@ -302,7 +305,6 @@ def on_pre_call(args): retained['x-retained'] = 'sent' ", json!({}), - json!({}), ); assert_eq!( wire.headers, @@ -314,52 +316,33 @@ def on_pre_call(args): } #[test] -fn post_call_receives_the_raw_response_and_the_payload_dicts_pre_call_saw() { +fn post_call_receives_the_raw_response_the_route_key_and_the_body_pre_call_saw() { before_send( c" def check(): - original_response, additional_args = logger.post + original_response, api_key, additional_args = logger.post assert original_response == 'raw response', original_response + assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) + assert additional_args == {'complete_input_dict': logger.pre['complete_input_dict']}, additional_args assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] - assert additional_args['headers'] is logger.pre['headers'] ", - json!({}), json!({"document": document(DOCUMENT)}), ); } -#[rstest] -#[case::every_phase_listens(c"{}", &["pre_call", "post_call"])] -#[case::no_input_callback( - c"{'input': False}", - &["_pre_call", "record_api_call_start_time", "record_post_call"] -)] -#[case::no_payload_consumer(c"{'payload': False}", &["record_api_call_start_time"])] -fn payload_callbacks_run_only_for_the_phases_someone_listens_to( - #[case] needed: &CStr, - #[case] expected_calls: &[&str], -) { - let script = std::ffi::CString::new(format!( - " -logger.needed = {needed} +#[test] +fn every_request_runs_the_full_pre_call_and_post_call() { + let wire = before_send( + c" def on_pre_call(args): args['complete_input_dict']['include_image_base64'] = True def check(): - assert logger.names() == {expected_calls:?}, logger.calls + assert logger.names() == ['pre_call', 'post_call'], logger.calls ", - needed = needed.to_str().unwrap(), - expected_calls = expected_calls, - )) - .unwrap(); - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(&script, json!({}), body.clone()); - let edited = json!({"document": document(DOCUMENT), "include_image_base64": true}); + json!({"document": document(DOCUMENT)}), + ); assert_eq!( wire.body, - if expected_calls.contains(&"pre_call") { - edited - } else { - body - } + json!({"document": document(DOCUMENT), "include_image_base64": true}) ); } diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy/tests/support.rs index 1663e11963e..444655ea77b 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/support.rs @@ -5,64 +5,89 @@ use pyo3::types::{PyDict, PyTuple}; use crate::{LegacyLogging, LegacySurface, PublicCall}; -/// Stand-ins for every litellm function the legacy contract calls. Tests share one -/// interpreter and run concurrently, so each stub is installed idempotently and forwards to -/// the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// The parameters of every `legacy_callbacks` function, as the real module declares them. +/// `tests/test_litellm/rust_bridge/test_legacy_callbacks.py` pins this file to the Python +/// signatures, and [`namespace`] binds every fake call against it. +pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + +/// Stand-ins for `legacy_callbacks`, the only Python module the crate calls. Tests +/// share one interpreter and run concurrently, so each fake is installed idempotently and +/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// Every fake is bound against the contract first, so a call the real module would reject +/// fails here too. const STUBS: &CStr = c" import contextvars +import inspect +import json import sys +import traceback import types -for name in ( - 'litellm', - 'litellm.utils', - 'litellm.types', - 'litellm.types.utils', - 'litellm._internal_context', - 'litellm.litellm_core_utils', - 'litellm.litellm_core_utils.logging_worker', - 'litellm.litellm_core_utils.litellm_logging', - 'litellm.rust_bridge', - 'litellm.rust_bridge.legacy_callbacks', -): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): sys.modules.setdefault(name, types.ModuleType(name)) legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] -legacy.setup = lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - bridge_owned=True, -) -legacy.deployment_callbacks_needed = lambda: True -legacy.check_limits = lambda arguments: arguments['logger'].check_limits(arguments) -legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True) -legacy.success_bookkeeping = lambda logger, response, start, end, asynchronous: logger.record( - 'success_bookkeeping', asynchronous -) -legacy.failure_bookkeeping = lambda logger, error, start, end, asynchronous: logger.record( - 'failure_bookkeeping', asynchronous -) -legacy.finalize = lambda response, logger, kwargs, start, end: logger.record('finalize', response) +CONTRACT = json.loads(python_contract) -utils = sys.modules['litellm.utils'] -utils.async_pre_call_deployment_hook = lambda kwargs, call_type: kwargs['logger'].hook( - 'pre', kwargs, call_type -) -utils.async_post_call_success_deployment_hook = lambda kwargs, response, call_type: kwargs[ - 'logger' -].hook('success', response, call_type) -utils.async_post_call_failure_deployment_hook = lambda kwargs, error, call_type: kwargs[ - 'logger' -].hook('failure', error, call_type) -utils._restore_correlation_context_if_supported = lambda logger: logger.record('restore', None) -internal = sys.modules['litellm._internal_context'] -if not hasattr(internal, 'is_internal_call'): - internal.is_internal_call = contextvars.ContextVar('is_internal_call', default=False) +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) -sys.modules['litellm.types.utils'].CustomPricingLiteLLMParams = type( - 'CustomPricingLiteLLMParams', (), {'model_fields': {'ocr_cost_per_page': None}} -) + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'credential_list': lambda: [], + 'warn_unknown_credential': lambda name, loaded: None, + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) unraisable = sys.modules.setdefault( @@ -77,20 +102,6 @@ def unraisable_from(owner): return [error for source, error in unraisable.events if source is owner] -class Worker: - def ensure_initialized_and_enqueue(self, coroutine): - return coroutine.enqueue() - - -class Executor: - def submit(self, run, handler, *args): - handler.__self__.record('submit', args) - - -sys.modules['litellm.litellm_core_utils.logging_worker'].GLOBAL_LOGGING_WORKER = Worker() -sys.modules['litellm.litellm_core_utils.litellm_logging'].executor = Executor() - - class StubCoroutine: def __init__(self, logger): self.logger = logger @@ -106,7 +117,6 @@ class StubCoroutine: class StubLogger: def __init__(self): self.calls = [] - self.needed = {} self.hooks = {} self.on_enqueue = lambda coroutine: None @@ -147,6 +157,7 @@ logger = StubLogger() /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); py.run(script, Some(&locals), Some(&locals)).unwrap(); locals diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 9b9d29108f6..3094d7b88d2 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue}; +use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; @@ -19,12 +19,16 @@ const TIMING: Timing = Timing { fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { LegacyLogging { - logger: Some(PythonLogger::new(local(locals, "logger").unbind(), true)), + logger: Some(PythonLogger::new(local(locals, "logger").unbind())), ..legacy_call(py, locals, asynchronous) } } -fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep { +fn succeed( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, +) -> LifecycleStep { let response = local(locals, "response").unbind(); logging .emit( @@ -35,7 +39,7 @@ fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLoggi .unwrap() } -fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep { +fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> LifecycleStep { let failure = PyErr::from_value(local(locals, "failure")); logging .emit( @@ -51,20 +55,14 @@ fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) #[rstest] #[case::sync_listened(false, c"", &["submit"])] -#[case::sync_unlistened(false, c"logger.needed = {'sync_success': False}", &["success_bookkeeping"])] #[case::async_listened( true, c"", &["async_success_handler", "enqueued", "sync_success_for_async_call"] )] -#[case::async_unlistened( - true, - c"logger.needed = {'async_success': False, 'sync_success_async': False}", - &["success_bookkeeping"] -)] #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] -fn success_reaches_only_the_callbacks_that_listen( +fn success_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, #[case] expected: &[&str], @@ -76,7 +74,7 @@ fn success_reaches_only_the_callbacks_that_listen( let mut logging = logged(py, &locals, asynchronous); assert!(matches!( succeed(py, &locals, &mut logging), - AdapterStep::Done + LifecycleStep::Done )); let names: Vec = local(&locals, "logger") .call_method0("names") @@ -109,7 +107,10 @@ fn internal_calls_skip_failure_callbacks_only_when_asynchronous( internal: true, ..logged(py, &locals, asynchronous) }; - assert!(matches!(fail(py, &locals, &mut logging), AdapterStep::Done)); + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Done + )); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -157,7 +158,7 @@ logger = FailingLogger() let mut logging = logged(py, &locals, true); assert!(matches!( succeed(py, &locals, &mut logging), - AdapterStep::Done + LifecycleStep::Done )); assert!( logging @@ -173,14 +174,8 @@ logger = FailingLogger() #[rstest] #[case::sync_listened(false, c"", &["failure_handler"])] -#[case::sync_unlistened(false, c"logger.needed = {'sync_failure': False}", &["failure_bookkeeping"])] #[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] -#[case::async_unlistened( - true, - c"logger.needed = {'sync_failure': False, 'async_failure': False}", - &["failure_bookkeeping", "failure_bookkeeping"] -)] -fn failure_reaches_only_the_callbacks_that_listen( +fn failure_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, #[case] expected: &[&str], @@ -192,7 +187,10 @@ fn failure_reaches_only_the_callbacks_that_listen( let mut logging = logged(py, &locals, asynchronous); let step = fail(py, &locals, &mut logging); let awaits_async_handler = expected.contains(&"async_failure_handler"); - assert_eq!(matches!(step, AdapterStep::Await(_)), awaits_async_handler); + assert_eq!( + matches!(step, LifecycleStep::Await(_)), + awaits_async_handler + ); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -227,7 +225,7 @@ logger = FailingLogger() let mut logging = logged(py, &locals, true); assert!(matches!( fail(py, &locals, &mut logging), - AdapterStep::Await(_) + LifecycleStep::Await(_) )); assert!( logging @@ -265,7 +263,7 @@ fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( }; let expected = result.as_ref().err().map(|error| error.value(py).clone()); match logging.resume(py, result) { - Ok(step) => assert!(done && matches!(step, AdapterStep::Done)), + Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), Err(propagated) => { assert!(!done); assert!(propagated.value(py).is(expected.unwrap())); diff --git a/litellm-rust/crates/callbacks/Cargo.toml b/litellm-rust/crates/callbacks/Cargo.toml index 4b966271478..a68ebc26a8d 100644 --- a/litellm-rust/crates/callbacks/Cargo.toml +++ b/litellm-rust/crates/callbacks/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-auth.workspace = true serde_json.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs index e6f88fd9709..3bf12e553b9 100644 --- a/litellm-rust/crates/callbacks/src/event.rs +++ b/litellm-rust/crates/callbacks/src/event.rs @@ -1,6 +1,6 @@ use std::time::{SystemTime, UNIX_EPOCH}; -use serde_json::{Map, Value}; +use serde_json::Value; /// Seconds since the Unix epoch, on one clock for every host. pub fn epoch_seconds() -> f64 { @@ -33,34 +33,10 @@ pub struct RequestContext { pub custom_llm_provider: String, /// The route's parameters before the provider transformation. pub optional_params: Value, - pub passthrough_fields: Passthrough, /// Optional-param names that carry credentials and must be redacted when logged. pub secret_fields: Vec, -} - -/// Body keys whose values are the caller's inputs, unchanged by the route. The only way to -/// build one is to compare the two, so a route cannot name a key it rewrote. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct Passthrough(Vec); - -impl Passthrough { - pub fn unchanged(caller: &Map, body: &Value) -> Self { - Self( - caller - .iter() - .filter(|(name, value)| body.get(name.as_str()) == Some(*value)) - .map(|(name, _)| name.clone()) - .collect(), - ) - } - - pub fn iter(&self) -> impl Iterator { - self.0.iter().map(String::as_str) - } - - pub fn contains(&self, name: &str) -> bool { - self.0.iter().any(|field| field == name) - } + /// The credential the route resolved for the provider call. + pub api_key: Option, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -78,6 +54,9 @@ pub enum FailureOrigin { #[derive(Clone, Debug, PartialEq)] pub enum CallEvent { + Started { + start_time: f64, + }, ResponseReceived { raw: RawResponse, }, @@ -89,47 +68,3 @@ pub enum CallEvent { origin: FailureOrigin, }, } - -#[cfg(test)] -mod tests { - use rstest::rstest; - use serde_json::json; - - use super::*; - - #[rstest] - #[case::unchanged_scalar(json!({"pages": [0]}), json!({"pages": [0]}), &["pages"])] - #[case::unchanged_explicit_null(json!({"pages": null}), json!({"pages": null}), &["pages"])] - #[case::unchanged_nested_object( - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}), - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}, "model": "m"}), - &["document"] - )] - #[case::rewritten_value( - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}), - json!({"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}), - &[] - )] - #[case::dropped_nested_field( - json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "document_name": "b.png"}}), - json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}), - &[] - )] - #[case::added_nested_field( - json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}), - json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "detail": "high"}}), - &[] - )] - #[case::reordered_array(json!({"pages": [0, 1]}), json!({"pages": [1, 0]}), &[])] - #[case::consumed_by_the_route(json!({"api_key": "k", "pages": [0]}), json!({"pages": [0]}), &["pages"])] - #[case::added_by_the_route(json!({}), json!({"model": "m"}), &[])] - #[case::non_object_body(json!({"pages": [0]}), json!([{"pages": [0]}]), &[])] - fn passthrough_is_exactly_the_callers_unchanged_keys( - #[case] caller: Value, - #[case] body: Value, - #[case] expected: &[&str], - ) { - let passthrough = Passthrough::unchanged(caller.as_object().unwrap(), &body); - assert_eq!(passthrough.iter().collect::>(), expected); - } -} diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/callbacks/src/run.rs index 57bf134f345..5accaa2e25e 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/callbacks/src/run.rs @@ -11,6 +11,7 @@ where H: Host, { let start_time = epoch_seconds(); + let _ = host.emit(&CallEvent::Started { start_time }).await; let mut result = None; let outcome = loop { let step = match machine.resume(result.take()).await { @@ -102,6 +103,7 @@ mod tests { async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { self.seen.lock().unwrap().push(match event { + CallEvent::Started { .. } => "started".into(), CallEvent::Succeeded { .. } => "succeeded".into(), CallEvent::Failed { .. } => "failed".into(), other => format!("{other:?}"), @@ -124,7 +126,7 @@ mod tests { assert_eq!(outcome, Ok(())); assert_eq!( *host.seen.lock().unwrap(), - ["route:project", "route:send", "succeeded"] + ["started", "route:project", "route:send", "succeeded"] ); } @@ -133,7 +135,7 @@ mod tests { let host = Recording::default(); let outcome = run(scripted(&[], Err("boom")), &host).await; assert_eq!(outcome, Err("boom")); - assert_eq!(*host.seen.lock().unwrap(), ["failed"]); + assert_eq!(*host.seen.lock().unwrap(), ["started", "failed"]); let host = Recording { fail: Some("send"), @@ -143,7 +145,35 @@ mod tests { assert_eq!(outcome, Err("host failed")); assert_eq!( *host.seen.lock().unwrap(), - ["route:project", "route:send", "failed"] + ["started", "route:project", "route:send", "failed"] ); } + + struct StartTimes(Mutex>); + + impl Host for StartTimes { + async fn route(&self, _: &'static str) -> Result<(), &'static str> { + Ok(()) + } + + async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { + if let CallEvent::Started { start_time } + | CallEvent::Succeeded { + timing: Timing { start_time, .. }, + } = event + { + self.0.lock().unwrap().push(*start_time); + } + Err("observer failed") + } + } + + #[tokio::test] + async fn started_opens_the_call_at_the_terminal_start_time_and_cannot_fail_it() { + let host = StartTimes(Mutex::default()); + assert_eq!(run(scripted(&["project"], Ok(())), &host).await, Ok(())); + let times = host.0.lock().unwrap(); + assert_eq!(times.len(), 2); + assert_eq!(times[0], times[1]); + } } diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 33cb8a8d32a..a6af190eb91 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,5 +1,6 @@ use futures_util::future::BoxFuture; -use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest}; +use litellm_auth::SecretValue; +use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, @@ -36,6 +37,7 @@ pub(crate) struct OcrCallHooks { custom_llm_provider: &'static str, optional_params: Value, secret_fields: Vec, + api_key: Option, } impl OcrCallHooks { @@ -51,22 +53,19 @@ impl OcrCallHooks { .filter(|name| is_secret_param(name)) .cloned() .collect(), + api_key: request.connection.api_key.clone(), } } } impl CallHooks for OcrCallHooks { - fn before_send( - &self, - wire: WireRequest, - passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result> { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { let context = RequestContext { model: self.model.clone(), custom_llm_provider: self.custom_llm_provider.into(), optional_params: self.optional_params.clone(), - passthrough_fields, secret_fields: self.secret_fields.clone(), + api_key: self.api_key.clone(), }; Box::pin(self.host.before_send(wire, context)) } diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index e7f77acc3f8..c977f721a70 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -21,8 +21,8 @@ mod cohere_tests; #[path = "../../tests/deepseek_ocr.rs"] mod deepseek_tests; #[cfg(test)] -#[path = "../../tests/ocr/passthrough.rs"] -mod passthrough_tests; +#[path = "../../tests/ocr/document.rs"] +mod document_tests; #[cfg(test)] #[path = "../../tests/reducto_ocr.rs"] mod reducto_tests; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 24c3f43e2b4..8ac038290b7 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,4 +1,4 @@ -use litellm_auth::{InputSource, Sourced}; +use litellm_auth::{InputSource, SecretValue, Sourced}; use litellm_llms::base_llm::ocr::transformation::{ OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env, }; @@ -22,7 +22,7 @@ pub(crate) fn prepare_request( .config .get_api_key_env_var() .and_then(credential_env) - .map(|value| Sourced::new(value, InputSource::Environment)) + .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) }) }); let dynamic_api_base = credentials.dynamic_api_base.or_else(|| { diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index d12b8cfee95..14b34ea4564 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -277,12 +277,18 @@ mod tests { #[test] fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { - api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("explicit-key"), + InputSource::Deployment, + )), api_base: Some(Sourced::new( "https://explicit.test".into(), InputSource::Deployment, )), - dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("dynamic-key"), + InputSource::Environment, + )), dynamic_api_base: Some(Sourced::new( "https://dynamic.test".into(), InputSource::Request, @@ -292,7 +298,7 @@ mod tests { connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), Some("dynamic-key") ); assert_eq!( @@ -318,22 +324,31 @@ mod tests { fn empty_or_missing_dynamic_credentials_preserve_explicit_values( #[case] dynamic_value: Option<&str>, ) { - let dynamic = + let dynamic_key = dynamic_value.map(|value| { + Sourced::new( + litellm_auth::SecretValue::new(value), + InputSource::Environment, + ) + }); + let dynamic_base = dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { - api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("explicit-key"), + InputSource::Deployment, + )), api_base: Some(Sourced::new( "https://explicit.test".into(), InputSource::Deployment, )), - dynamic_api_key: dynamic.clone(), - dynamic_api_base: dynamic, + dynamic_api_key: dynamic_key, + dynamic_api_base: dynamic_base, }); assert_eq!( connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), Some("explicit-key") ); assert_eq!( @@ -356,11 +371,18 @@ mod tests { ) { let connection = OcrConfigKind::AzureDocumentIntelligence.resolve_connection_params( OcrCredentialInputs { - api_key: explicit_key - .map(|value| Sourced::new(value.into(), InputSource::Deployment)), + api_key: explicit_key.map(|value| { + Sourced::new( + litellm_auth::SecretValue::new(value), + InputSource::Deployment, + ) + }), api_base: explicit_base .map(|value| Sourced::new(value.into(), InputSource::Deployment)), - dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("dynamic-key"), + InputSource::Environment, + )), dynamic_api_base: Some(Sourced::new( "https://dynamic.test".into(), InputSource::Deployment, @@ -371,7 +393,7 @@ mod tests { connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), explicit_key.map(|_| "dynamic-key") ); assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 75202ed52a5..6316088dec8 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, path::PathBuf, time::Duration}; use bytes::Bytes; -use litellm_auth::{InputSource, TokenProviderHandle}; +use litellm_auth::{InputSource, SecretValue, TokenProviderHandle}; use litellm_core_utils::call_arguments::CallArguments; use litellm_llms::base_llm::ocr::{ error::Error, @@ -56,7 +56,7 @@ pub struct OcrFileContent { /// credentials, and per-field provenance in `input_sources`. #[derive(Clone, Debug, Default)] pub struct OcrConnectionInputs { - pub api_key: Option, + pub api_key: Option, pub api_base: Option, pub extra_headers: Map, pub timeout: Option, @@ -237,6 +237,16 @@ mod tests { .unwrap() } + #[test] + fn connection_inputs_debug_hides_the_api_key() { + let inputs = OcrConnectionInputs { + api_key: Some(SecretValue::new("caller-api-key")), + ..OcrConnectionInputs::default() + }; + + assert!(!format!("{inputs:?}").contains("caller-api-key")); + } + #[test] fn from_inputs_applies_connection_overrides_with_field_sources() { let request = LiteLLMOcrRequest::from_inputs( @@ -245,7 +255,7 @@ mod tests { None, Default::default(), OcrConnectionInputs { - api_key: Some(" key ".into()), + api_key: Some(SecretValue::new(" key ")), api_base: Some("".into()), extra_headers: json!({"x-a": "1"}).as_object().unwrap().clone(), timeout: Some(Duration::from_secs(7)), @@ -259,7 +269,7 @@ mod tests { .unwrap(); let api_key = request.credentials.api_key.as_ref().unwrap(); - assert_eq!(api_key.clone().into_value(), "key"); + assert_eq!(api_key.value().expose(), "key"); assert_eq!(api_key.source(), InputSource::Request); assert!(request.credentials.api_base.is_none()); assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 29345e38885..b9c60f57e3c 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, time::Duration}; -use litellm_auth::InputSource; +use litellm_auth::{InputSource, SecretValue}; use litellm_llms::base_llm::ocr::{ error::Error, transformation::{OcrDocument, decode_request_value}, @@ -44,7 +44,7 @@ pub fn consumed_optional_param_names( pub struct OcrWireRequest { pub model: String, pub document: D, - pub api_key: Option, + pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, pub extra_headers: Option>, diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 01a4e5efb3b..1fc4d6c2b9e 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -69,7 +69,7 @@ async fn rejects_invalid_pages_features_and_format( let result = decode_request(OcrWireRequest { model: "azure_ai/doc-intelligence/prebuilt-read".into(), document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: Some(base), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 1f591d74d5d..ca2e14a7f0d 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -81,7 +81,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { let request = OcrWireRequest { model: "mistral/model".into(), document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: None, custom_llm_provider: None, extra_headers: None, @@ -97,7 +97,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { decode_request(OcrWireRequest { model: "model".into(), document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: None, custom_llm_provider: Some("unknown".into()), extra_headers: None, @@ -194,6 +194,7 @@ async fn facade_uses_the_injected_http_client() { fn event_name(event: &CallEvent) -> &'static str { match event { + CallEvent::Started { .. } => "started", CallEvent::ResponseReceived { .. } => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", @@ -235,7 +236,7 @@ async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { } #[tokio::test] -async fn before_send_context_names_passthrough_fields_and_secrets() { +async fn before_send_context_names_the_route_and_its_secrets() { let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; let observed = Arc::new(Mutex::new(None)); let captured = observed.clone(); @@ -254,8 +255,6 @@ async fn before_send_context_names_passthrough_fields_and_secrets() { assert_eq!(context.custom_llm_provider, "mistral"); assert_eq!(context.model, "model"); assert_eq!(wire.body["pages"], json!([0])); - assert!(context.passthrough_fields.contains("pages")); - assert!(context.passthrough_fields.contains("document")); assert!(context.secret_fields.is_empty()); assert_eq!(context.optional_params["req_format"], "native"); @@ -279,7 +278,6 @@ async fn before_send_context_names_passthrough_fields_and_secrets() { perform_ocr_with(host).await.unwrap(); server.await.unwrap(); let context = observed.lock().unwrap().take().unwrap(); - assert!(!context.passthrough_fields.contains("document")); assert_eq!(context.secret_fields, ["client_secret"]); } @@ -296,7 +294,7 @@ async fn lifecycle_orders_hooks_and_emits_one_success() { server.await.unwrap(); assert_eq!( *events.lock().unwrap(), - ["before_send", "response", "success"] + ["started", "before_send", "response", "success"] ); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -311,7 +309,10 @@ async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { ); let error = perform_ocr_with(host).await.unwrap_err(); assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); } #[tokio::test] @@ -330,7 +331,10 @@ async fn upstream_failure_emits_one_terminal_failure() { ); assert!(perform_ocr_with(host).await.is_err()); server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); assert_eq!(seen.lock().unwrap().len(), 1); } diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs new file mode 100644 index 00000000000..5e10ce3ad2a --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/document.rs @@ -0,0 +1,152 @@ +use litellm_callbacks::event::WireRequest; +use litellm_llms::base_llm::ocr::error::Error; +use rstest::rstest; +use serde_json::{Value, json}; + +use super::test_support::{ + MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, + wire_request_with_document, +}; +use crate::ocr::route::LocalOcrHost; + +#[derive(Clone, Copy, Debug)] +enum Route { + Mistral, + AzureAi, + VertexMistral, + AzureCohereParse, + Cohere, +} + +impl Route { + fn model(self) -> &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } +} + +/// What the host does to the wire request in `before_send`. +#[derive(Clone, Copy, Debug)] +enum Host { + Detached, + ReplacesDocument, +} + +const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + +impl Host { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::Detached => (name, value), + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } +} + +struct Sent { + result: Result<(), Error>, + provider_body: Option, +} + +async fn send(route: Route, host: Host, document_base: &str) -> Sent { + let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; + let document_type = route.document_type(); + let document = + json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); + let request = wire_request_with_document(route.model(), &base, document, route.options()); + let local = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); + let result = perform_ocr_with(local).await.map(|_| ()); + match result { + Ok(()) => provider.await.unwrap(), + Err(_) => provider.abort(), + } + let provider_body = seen + .lock() + .unwrap() + .first() + .map(|request| request_body(request)); + Sent { + result, + provider_body, + } +} + +fn served_document_uri() -> String { + use base64::Engine; + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ) +} + +#[rstest] +#[case::azure_ai(Route::AzureAi)] +#[case::vertex_mistral(Route::VertexMistral)] +#[case::azure_cohere_parse(Route::AzureCohereParse)] +#[tokio::test] +async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::Detached, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(served_document_uri()) + ); +} + +#[rstest] +#[tokio::test] +async fn document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, +) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::ReplacesDocument, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(REPLACED_DOCUMENT) + ); +} diff --git a/litellm-rust/crates/core/tests/ocr/passthrough.rs b/litellm-rust/crates/core/tests/ocr/passthrough.rs deleted file mode 100644 index 0273b48664d..00000000000 --- a/litellm-rust/crates/core/tests/ocr/passthrough.rs +++ /dev/null @@ -1,282 +0,0 @@ -use std::{ - collections::BTreeSet, - sync::{Arc, Mutex}, -}; - -use litellm_callbacks::event::{RequestContext, WireRequest}; -use litellm_llms::base_llm::ocr::error::Error; -use rstest::rstest; -use rstest_reuse::{self, apply, template}; -use serde_json::{Map, Value, json}; - -use super::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, - wire_request_with_document, -}; -use crate::ocr::route::LocalOcrHost; - -#[derive(Clone, Copy, Debug)] -enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, -} - -impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } -} - -#[derive(Clone, Copy, Debug)] -enum Source { - Inline, - Remote, - RemoteWithExtraField, -} - -/// What the host does to the wire request in `before_send`. -#[derive(Clone, Copy, Debug)] -enum Host { - Detached, - /// What `litellm-callbacks-legacy` does before `pre_call`: every passthrough body key - /// is replaced by the caller's own value. - Realiasing, - ReplacesDocument, -} - -const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - -impl Host { - fn before_send( - self, - caller: &Map, - wire: WireRequest, - context: &RequestContext, - ) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::Realiasing => { - let aliased = context - .passthrough_fields - .contains(&name) - .then(|| caller.get(&name).cloned()) - .flatten() - .unwrap_or(value); - (name, aliased) - } - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } -} - -struct Sent { - caller: Map, - result: Result<(), Error>, - before_send: Option<(WireRequest, RequestContext)>, - provider_body: Option, -} - -fn caller_document(route: Route, source: Source, document_base: &str) -> Value { - let document_type = route.document_type(); - let remote = format!("{document_base}/scan.png"); - match source { - Source::Inline => { - json!({"type": document_type, document_type: "data:image/png;base64,YWJj"}) - } - Source::Remote => json!({"type": document_type, document_type: remote}), - Source::RemoteWithExtraField => { - json!({"type": document_type, document_type: remote, "document_name": "scan.png"}) - } - } -} - -async fn send(route: Route, source: Source, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document = caller_document(route, source, document_base); - let caller: Map = route - .options() - .as_object() - .unwrap() - .clone() - .into_iter() - .chain([("document".to_string(), document.clone())]) - .collect(); - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host_caller = caller.clone(); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(host.before_send(&host_caller, wire, context)) - }); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - let before_send = observed.lock().unwrap().take(); - Sent { - caller, - result, - before_send, - provider_body, - } -} - -fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) -} - -#[template] -#[rstest] -fn every_route_and_source( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, - #[values(Source::Inline, Source::Remote, Source::RemoteWithExtraField)] source: Source, -) { -} - -#[template] -#[rstest] -fn every_route( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, -) { -} - -#[template] -#[rstest] -#[case::azure_ai(Route::AzureAi)] -#[case::vertex_mistral(Route::VertexMistral)] -#[case::azure_cohere_parse(Route::AzureCohereParse)] -fn inlining_routes(#[case] route: Route) {} - -#[apply(every_route_and_source)] -#[tokio::test] -async fn passthrough_fields_are_exactly_the_caller_values_sent_unchanged( - route: Route, - source: Source, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, source, Host::Detached, &document_base).await; - sent.result.unwrap(); - let (wire, context) = sent.before_send.unwrap(); - let passthrough: BTreeSet<&str> = context.passthrough_fields.iter().collect(); - let unchanged: BTreeSet<&str> = sent - .caller - .iter() - .filter(|(name, value)| wire.body.get(name.as_str()) == Some(*value)) - .map(|(name, _)| name.as_str()) - .collect(); - assert_eq!( - passthrough, - unchanged, - "body: {:#}\ncaller: {:#}", - wire.body, - Value::Object(sent.caller.clone()) - ); -} - -#[apply(every_route_and_source)] -#[tokio::test] -async fn realiasing_leaves_the_provider_request_unchanged(route: Route, source: Source) { - let (document_base, _documents) = document_server().await; - let detached = send(route, source, Host::Detached, &document_base).await; - let realiased = send(route, source, Host::Realiasing, &document_base).await; - detached.result.unwrap(); - realiased.result.unwrap(); - assert_eq!(realiased.provider_body, detached.provider_body); -} - -#[apply(inlining_routes)] -#[tokio::test] -async fn inlining_routes_send_the_downloaded_document( - route: Route, - #[values(Host::Detached, Host::Realiasing)] host: Host, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Source::Remote, host, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); -} - -#[apply(every_route)] -#[tokio::test] -async fn document_replaced_by_the_host_reaches_the_provider(route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send( - route, - Source::Remote, - Host::ReplacesDocument, - &document_base, - ) - .await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); -} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index f3adf27cfa6..44313d5f552 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use futures_util::future::BoxFuture; -use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_callbacks::event::WireRequest; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::{CallHooks, OcrClient}, @@ -23,11 +23,7 @@ use crate::ocr::{ pub(crate) struct NoHooks; impl CallHooks for NoHooks { - fn before_send( - &self, - wire: WireRequest, - _passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result> { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { Box::pin(async move { Ok(wire) }) } @@ -70,7 +66,7 @@ pub(crate) fn wire_request_with_document( decode_request(OcrWireRequest { model: model.into(), document, - api_key: Some("test-key".into()), + api_key: Some(litellm_auth::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index a3fdd2340b3..903ccb06c77 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -1,9 +1,10 @@ - Target invariants; implementation and runtime validation may lag these rules -- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `CallbackAdapter`/`RouteHost` traits +- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`RouteHost` traits - No LiteLLM domain dependencies beyond `litellm-callbacks`: no route types, no `Logging` policy, no public API registration, no cdylib build features - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business - `RouteHost::invoke` receives the keyword view the adapter's `begin` returned, not the caller's dict; a route host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) - - A failure that surfaces inside the call, including a host op the call asked for, is mapped through the route's `map_failure`; a failure in `begin` or `after_success` is raised as is + - A native failure, including one a host op returns as `HostOpError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is + - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads - Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized` diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index f1bc3142a25..4aa7a2163ca 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -11,7 +11,7 @@ pub fn missing_state() -> PyErr { /// What an adapter step produced: either the value the driver asked for, or a Python /// awaitable the driver hands back to the caller's task before asking again. -pub enum AdapterStep { +pub enum LifecycleStep { Await(Py), Arguments(Py), Wire(Box), @@ -28,55 +28,74 @@ pub enum PublicValue<'a> { /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in /// order: `begin` before the machine starts, `before_send` and `emit` while it runs, /// `after_success` and one terminal `emit` after it completes. Whenever a step returns -/// [`AdapterStep::Await`], the driver awaits it in the caller's task and continues the +/// [`LifecycleStep::Await`], the driver awaits it in the caller's task and continues the /// same step through `resume`. /// /// A step that fails with an ordinary exception fails the call with that exception, /// except on a terminal event, where the adapter is expected to report and swallow its /// own errors. An exception that is not a `PyException`, such as a cancellation, ends /// the call without further dispatch. -pub trait CallbackAdapter: Send + Sync { +pub trait PythonLifecycle: Send + Sync { fn begin( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult; + ) -> PyResult; fn before_send( &mut self, py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult; + ) -> PyResult; fn after_success( &mut self, py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult; + ) -> PyResult; fn emit( &mut self, py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult; + ) -> PyResult; - fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; fn close(&mut self, py: Python<'_>); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; } -/// The Python side of one route: answers the route's own operations, builds the public -/// response and maps failures to public exceptions. -pub trait RouteHost: Send + Sync { - type Route: Route; +/// Why a route operation the host answered did not produce a result: the route's own code +/// rejected it, which the route classifies like any other native failure, or Python code +/// raised, which reaches the caller as it was raised. +#[derive(Debug)] +pub enum HostOpError { + Native(E), + Python(PyErr), +} - /// `arguments` is the keyword view the callback adapter's `begin` produced, not the +impl From for HostOpError { + fn from(error: PyErr) -> Self { + Self::Python(error) + } +} + +/// The Python side of one route: answers the route's own operations, builds the public +/// response and classifies native failures into public exceptions. +pub trait RouteHost: Send + Sync { + type Route: Route; + + /// The public exception a native failure maps to, kept as a value until the driver + /// raises it. + type Failure: Into; + + /// `arguments` is the keyword view the lifecycle's `begin` produced, not the /// caller's own dict. A route host that projects from it inherits whatever that /// adapter rewrote. fn invoke( @@ -84,7 +103,7 @@ pub trait RouteHost: Send + Sync { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: ::Op, - ) -> PyResult<::OpResult>; + ) -> Result<::OpResult, HostOpError<::Error>>; fn complete( &mut self, @@ -92,12 +111,14 @@ pub trait RouteHost: Send + Sync { response: ::Response, ) -> PyResult>; - fn native_error(error: ::Error) -> PyErr; + fn classify( + &self, + py: Python<'_>, + error: ::Error, + ) -> PyResult; fn host_error(error: &PyErr) -> ::Error; - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult; - fn close(&mut self, py: Python<'_>); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; diff --git a/litellm-rust/crates/host-python/src/argument.rs b/litellm-rust/crates/host-python/src/argument.rs new file mode 100644 index 00000000000..34e07cdfbd5 --- /dev/null +++ b/litellm-rust/crates/host-python/src/argument.rs @@ -0,0 +1,51 @@ +use pyo3::{prelude::*, types::PyDict}; + +/// The caller's own object for a public argument: the keyword if given, even an explicit +/// `None`, else the bound request's attribute. Every reader of a public Python call uses +/// this rule, so the callbacks and the provider see one object per argument. +pub fn lookup<'py>( + kwargs: &Bound<'py, PyDict>, + request: &Bound<'py, PyAny>, + name: &str, +) -> PyResult>> { + if let Some(value) = kwargs.get_item(name)? { + return Ok(Some(value)); + } + request.getattr_opt(name) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { + crate::initialize_python(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +key = object() +document = {'type': 'document_url'} +class Request: + api_key = 'from-request' + api_base = 'from-request' + document = document +request = Request() +kwargs = {'api_key': key, 'api_base': None} +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let item = |name: &str| locals.get_item(name).unwrap().unwrap(); + let kwargs = item("kwargs").cast_into::().unwrap(); + let request = item("request"); + let find = |name: &str| lookup(&kwargs, &request, name).unwrap(); + assert!(find("api_key").unwrap().is(item("key"))); + assert!(find("api_base").unwrap().is_none()); + assert!(find("document").unwrap().is(item("document"))); + assert!(find("model").is_none()); + }); + } +} diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 8bda13b44d0..7895d803cb5 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -12,7 +12,9 @@ use pyo3::prelude::*; use pyo3::types::PyDict; use tokio::sync::Mutex; -use crate::adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state}; +use crate::adapter::{ + HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, +}; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -43,6 +45,7 @@ enum Stage { #[derive(Clone, Copy)] enum Expect { + Started, Arguments, Wire, Emitted, @@ -66,7 +69,7 @@ where M: Machine> + 'static, { route: H, - adapter: Box, + adapter: Box, machine: Option>>>, arguments: Option>, started_at: f64, @@ -84,7 +87,7 @@ pub fn run_call( py: Python<'_>, machine: M, route: H, - adapter: Box, + adapter: Box, arguments: Py, asynchronous: bool, ) -> PyResult> @@ -146,9 +149,11 @@ where match (self.pending.take(), result) { (None, None) => { self.started_at = epoch_seconds(); - let arguments = self.arguments.take().ok_or_else(missing_state)?; - match self.adapter.begin(py, arguments, self.started_at) { - Ok(step) => self.on_adapter(py, step, Expect::Arguments), + let started = CallEvent::Started { + start_time: self.started_at, + }; + match self.adapter.emit(py, &started, None) { + Ok(step) => self.on_adapter(py, step, Expect::Started), Err(error) => self.adapter_failed(py, error), } } @@ -170,27 +175,28 @@ where fn on_adapter( &mut self, py: Python<'_>, - step: AdapterStep, + step: LifecycleStep, expect: Expect, ) -> PyResult { match (expect, step) { - (_, AdapterStep::Await(awaitable)) => { + (_, LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(expect)); Ok(ExecutionStep::Await(awaitable)) } - (Expect::Arguments, AdapterStep::Arguments(arguments)) => { + (Expect::Started, LifecycleStep::Done) => self.begin(py), + (Expect::Arguments, LifecycleStep::Arguments(arguments)) => { self.arguments = Some(arguments); self.stage = Stage::Call; self.resume_machine(py, None) } - (Expect::Wire, AdapterStep::Wire(wire)) => { + (Expect::Wire, LifecycleStep::Wire(wire)) => { self.resume_machine(py, Some(Ok(HostResult::BeforeSend(wire)))) } - (Expect::Emitted, AdapterStep::Done) => { + (Expect::Emitted, LifecycleStep::Done) => { self.resume_machine(py, Some(Ok(HostResult::Emitted))) } - (Expect::Response, AdapterStep::Response(response)) => self.succeeded(py, response), - (Expect::Terminal, AdapterStep::Done) => match &self.stage { + (Expect::Response, LifecycleStep::Response(response)) => self.succeeded(py, response), + (Expect::Terminal, LifecycleStep::Done) => match &self.stage { Stage::Succeeded(response) => Ok(ExecutionStep::Return(response.clone_ref(py))), Stage::Failed(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())), _ => Err(missing_state()), @@ -199,6 +205,14 @@ where } } + fn begin(&mut self, py: Python<'_>) -> PyResult { + let arguments = self.arguments.take().ok_or_else(missing_state)?; + match self.adapter.begin(py, arguments, self.started_at) { + Ok(step) => self.on_adapter(py, step, Expect::Arguments), + Err(error) => self.adapter_failed(py, error), + } + } + fn adapter_failed(&mut self, py: Python<'_>, error: PyErr) -> PyResult { match self.stage { Stage::Begin | Stage::AfterSuccess => self.failure(py, error, FailureOrigin::Host), @@ -248,14 +262,20 @@ where let answer = match op { HostOp::Route(op) => { let arguments = self.arguments.as_ref().ok_or_else(missing_state)?; - self.route - .invoke(py, arguments.bind(py), op) - .map(HostResult::Route) + match self.route.invoke(py, arguments.bind(py), op) { + Ok(result) => Ok(HostResult::Route(result)), + Err(HostOpError::Native(error)) => { + return self + .resume_core(py, Some(Err(HostFailure::Error(error)))) + .map(Next::Continue); + } + Err(HostOpError::Python(error)) => Err(error), + } } HostOp::BeforeSend { wire, context } => { match self.adapter.before_send(py, wire, &context) { - Ok(AdapterStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), - Ok(AdapterStep::Await(awaitable)) => { + Ok(LifecycleStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), + Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Wire)); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } @@ -264,8 +284,8 @@ where } } HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { - Ok(AdapterStep::Done) => Ok(HostResult::Emitted), - Ok(AdapterStep::Await(awaitable)) => { + Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), + Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Emitted)); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } @@ -360,11 +380,29 @@ where self.ended_at.get_or_insert_with(epoch_seconds); let error = match self.interrupted.take() { Some(retained) => PyErr::from_value(retained.into_bound(py).into_any()), - None => H::native_error(error), + None => self.classified(py, error), }; self.failure(py, error, FailureOrigin::Call) } + /// The route's public exception for a native failure. When classification itself + /// fails, that failure is raised with the native error's text as its `__context__`. + fn classified(&self, py: Python<'_>, error: ErrorOf) -> PyErr { + let native = error.to_string(); + let classifier_error = match self.route.classify(py, error) { + Ok(failure) => return failure.into(), + Err(classifier_error) => classifier_error, + }; + let attached = classifier_error.value(py).setattr( + "__context__", + PyRuntimeError::new_err(native).into_value(py), + ); + match attached { + Ok(()) => classifier_error, + Err(error) => error, + } + } + fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { let event = CallEvent::Succeeded { timing: self.timing(), @@ -386,18 +424,14 @@ where if is_cancellation(py, &error) { return Err(error); } - let public = match origin { - FailureOrigin::Call => self.route.map_failure(py, &error).unwrap_or(error), - FailureOrigin::Host => error, - }; let event = CallEvent::Failed { timing: self.timing(), origin, }; let step = self .adapter - .emit(py, &event, Some(PublicValue::Error(&public)))?; - self.stage = Stage::Failed(public.into_value(py)); + .emit(py, &event, Some(PublicValue::Error(&error)))?; + self.stage = Stage::Failed(error.into_value(py)); self.on_adapter(py, step, Expect::Terminal) } @@ -489,6 +523,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri #[derive(Clone, Debug, PartialEq, Eq)] struct Error(String); + impl std::fmt::Display for Error { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } + } + struct Synthetic; impl Route for Synthetic { @@ -518,8 +558,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri model: "model".into(), custom_llm_provider: "provider".into(), optional_params: serde_json::json!({}), - passthrough_fields: Default::default(), secret_fields: Vec::new(), + api_key: None, } } @@ -566,25 +606,46 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } + #[derive(Clone, Copy)] + enum OpScript { + Answer, + RaisePython, + RejectNatively, + } + struct SyntheticHost { log: Log, - fail_op: bool, + op: OpScript, + classifier_fails: bool, + } + + /// The fake route's public exception, kept as a value so a test sees what `classify` + /// produced before the driver raises it. + #[derive(Debug, PartialEq, Eq)] + struct Classified(String); + + impl From for PyErr { + fn from(classified: Classified) -> Self { + PyValueError::new_err(format!("classified: {}", classified.0)) + } } impl RouteHost for SyntheticHost { type Route = Synthetic; + type Failure = Classified; fn invoke( &mut self, _: Python<'_>, arguments: &Bound<'_, PyDict>, op: &'static str, - ) -> PyResult { + ) -> Result> { self.log.push(format!("route:{op}")); - if self.fail_op { - return Err(PyValueError::new_err("op failed")); + match self.op { + OpScript::Answer => Ok(format!("{op}:{}", arguments.len())), + OpScript::RaisePython => Err(PyValueError::new_err("op failed").into()), + OpScript::RejectNatively => Err(HostOpError::Native(Error("op rejected".into()))), } - Ok(format!("{op}:{}", arguments.len())) } fn complete(&mut self, py: Python<'_>, response: String) -> PyResult> { @@ -594,22 +655,18 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .unbind()) } - fn native_error(error: Error) -> PyErr { - PyValueError::new_err(error.0) + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + self.log.push(format!("classify:{error}")); + if self.classifier_fails { + return Err(pyo3::exceptions::PyTypeError::new_err("classifier failed")); + } + Ok(Classified(error.0)) } fn host_error(error: &PyErr) -> Error { Error(error.to_string()) } - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult { - self.log.push("map_failure"); - Ok(PyValueError::new_err(format!( - "mapped: {}", - error.value(py) - ))) - } - fn close(&mut self, _: Python<'_>) { self.log.push("route.close"); } @@ -632,13 +689,18 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri script: AdapterScript, } - impl CallbackAdapter for SyntheticAdapter { - fn begin(&mut self, _: Python<'_>, arguments: Py, _: f64) -> PyResult { + impl PythonLifecycle for SyntheticAdapter { + fn begin( + &mut self, + _: Python<'_>, + arguments: Py, + _: f64, + ) -> PyResult { self.log.push("begin"); if matches!(self.script, AdapterScript::FailBegin) { return Err(PyValueError::new_err("begin failed")); } - Ok(AdapterStep::Arguments(arguments)) + Ok(LifecycleStep::Arguments(arguments)) } fn before_send( @@ -646,9 +708,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri _: Python<'_>, wire: Box, _: &RequestContext, - ) -> PyResult { + ) -> PyResult { self.log.push("before_send"); - Ok(AdapterStep::Wire(Box::new(WireRequest { + Ok(LifecycleStep::Wire(Box::new(WireRequest { url: "rewritten".into(), ..*wire }))) @@ -659,17 +721,17 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, response: Py, _: Timing, - ) -> PyResult { + ) -> PyResult { self.log.push("after_success"); match self.script { - AdapterScript::ReplaceResponse => Ok(AdapterStep::Response( + AdapterScript::ReplaceResponse => Ok(LifecycleStep::Response( "replaced".into_pyobject(py)?.into_any().unbind(), )), AdapterScript::FailAfterSuccess => { Err(PyValueError::new_err("after_success failed")) } AdapterScript::Plain | AdapterScript::FailBegin => { - Ok(AdapterStep::Response(response)) + Ok(LifecycleStep::Response(response)) } } } @@ -679,8 +741,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult { + ) -> PyResult { self.log.push(match (event, public) { + (CallEvent::Started { .. }, None) => "started".into(), (CallEvent::ResponseReceived { raw }, None) => format!("response:{}", raw.body), (CallEvent::Succeeded { .. }, Some(PublicValue::Response(value))) => { format!("succeeded:{}", value.bind(py)) @@ -690,10 +753,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } _ => "unexpected".into(), }); - Ok(AdapterStep::Done) + Ok(LifecycleStep::Done) } - fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { + fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { Err(missing_state()) } @@ -709,15 +772,31 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn run_scripted( py: Python<'_>, machine: ScriptedMachine, - fail_op: bool, + op: OpScript, script: AdapterScript, asynchronous: bool, ) -> (PyResult>, Vec) { - let log = Log::default(); - let route = SyntheticHost { - log: Log(log.0.clone()), - fail_op, - }; + run_hosted( + py, + machine, + SyntheticHost { + log: Log::default(), + op, + classifier_fails: false, + }, + script, + asynchronous, + ) + } + + fn run_hosted( + py: Python<'_>, + machine: ScriptedMachine, + route: SyntheticHost, + script: AdapterScript, + asynchronous: bool, + ) -> (PyResult>, Vec) { + let log = Log(route.log.0.clone()); let adapter = SyntheticAdapter { log: Log(log.0.clone()), script, @@ -777,7 +856,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::Plain, asynchronous, ); @@ -785,6 +864,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert_eq!( log, [ + "started", "begin", "route:project", "before_send", @@ -800,28 +880,75 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri }); } + fn failing_machine() -> ScriptedMachine { + ScriptedMachine { + ops: vec![HostOp::Route("project")], + outcome: Some(Err(Error("provider exploded".into()))), + answers: Vec::new(), + } + } + #[test] - fn machine_failures_are_mapped_and_dispatched_once_as_call_failures() { + fn a_native_failure_is_classified_once_and_reported_classified() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); crate::initialize_python(); Python::attach(|py| { - let machine = ScriptedMachine { - ops: vec![HostOp::Route("project")], - outcome: Some(Err(Error("provider exploded".into()))), - answers: Vec::new(), - }; - let (result, log) = run_scripted(py, machine, false, AdapterScript::Plain, false); - let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "mapped: provider exploded"); + install_lifecycle_module(py); + for asynchronous in [false, true] { + let (result, log) = run_scripted( + py, + failing_machine(), + OpScript::Answer, + AdapterScript::Plain, + asynchronous, + ); + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "classified: provider exploded"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "classify:provider exploded", + "failed:Call:classified: provider exploded", + "adapter.close", + "route.close", + ] + ); + } + }); + } + + #[test] + fn a_native_rejection_from_a_host_operation_is_classified_once() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::RejectNatively, + AdapterScript::Plain, + false, + ); + assert_eq!( + result.unwrap_err().value(py).to_string(), + "classified: op rejected" + ); assert_eq!( log, [ + "started", "begin", "route:project", - "map_failure", - "failed:Call:mapped: provider exploded", + "classify:op rejected", + "failed:Call:classified: op rejected", "adapter.close", "route.close", ] @@ -830,18 +957,72 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } #[test] - fn host_operation_failures_interrupt_the_call_and_keep_the_python_exception() { + fn a_python_exception_from_a_host_operation_is_reported_as_raised() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); crate::initialize_python(); Python::attach(|py| { - let (result, log) = - run_scripted(py, success_machine(), true, AdapterScript::Plain, false); + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::RaisePython, + AdapterScript::Plain, + false, + ); let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "mapped: op failed"); - assert!(!log.contains(&"before_send".to_string())); - assert!(log.contains(&"failed:Call:mapped: op failed".to_string())); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "op failed"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "failed:Call:op failed", + "adapter.close", + "route.close", + ] + ); + }); + } + + #[test] + fn a_failing_classifier_surfaces_with_the_native_error_as_context() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + let (result, log) = run_hosted( + py, + failing_machine(), + SyntheticHost { + log: Log::default(), + op: OpScript::Answer, + classifier_fails: true, + }, + AdapterScript::Plain, + false, + ); + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "classifier failed"); + let context = error.value(py).getattr("__context__").unwrap(); + assert!(context.is_instance_of::()); + assert_eq!(context.str().unwrap().to_string(), "provider exploded"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "classify:provider exploded", + "failed:Call:classifier failed", + "adapter.close", + "route.close", + ] + ); }); } @@ -855,7 +1036,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::FailBegin, false, ); @@ -864,6 +1045,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert_eq!( log, [ + "started", "begin", "failed:Host:begin failed", "adapter.close", @@ -885,7 +1067,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::ReplaceResponse, asynchronous, ); @@ -908,7 +1090,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::FailAfterSuccess, asynchronous, ); @@ -938,12 +1120,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri struct Cancelling(Log); impl RouteHost for Cancelling { type Route = Synthetic; + type Failure = Classified; fn invoke( &mut self, py: Python<'_>, _: &Bound<'_, PyDict>, _: &'static str, - ) -> PyResult { + ) -> Result> { self.0.push("route"); Err(PyErr::from_value( py.import("asyncio") @@ -952,21 +1135,19 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .unwrap() .call0() .unwrap(), - )) + ) + .into()) } fn complete(&mut self, _: Python<'_>, _: String) -> PyResult> { Err(missing_state()) } - fn native_error(error: Error) -> PyErr { - PyValueError::new_err(error.0) + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + self.0.push("classify"); + Ok(Classified(error.0)) } fn host_error(error: &PyErr) -> Error { Error(error.to_string()) } - fn map_failure(&self, _: Python<'_>, _: &PyErr) -> PyResult { - self.0.push("map_failure"); - Err(missing_state()) - } fn close(&mut self, _: Python<'_>) {} fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { Ok(()) @@ -988,7 +1169,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ) .unwrap_err(); assert!(!error.is_instance_of::(py)); - assert_eq!(log.entries(), ["begin", "route", "adapter.close"]); + assert_eq!( + log.entries(), + ["started", "begin", "route", "adapter.close"] + ); }); } diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index bb0b5b1c3b1..8889b1513db 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -1,9 +1,10 @@ //! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and //! asyncio glue, and the driver that runs a native [`Machine`](litellm_callbacks::machine::Machine) -//! against a Python route host and a callback adapter. Everything here is Python-specific by +//! against a Python route host and a Python lifecycle. Everything here is Python-specific by //! construction; another host language gets its own crate of the same shape. mod adapter; +mod argument; mod callable; mod driver; mod execution; @@ -11,7 +12,10 @@ mod gil; mod handle; mod marshal; -pub use adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state}; +pub use adapter::{ + HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, +}; +pub use argument::lookup; pub use callable::wrap_failure; pub use driver::run_call; pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value}; diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 87945bf8785..2e398d0287e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -150,7 +150,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { api_key: inputs.api_key.and_then(|key| { inputs .dynamic_api_key - .filter(|value| !value.value().is_empty()) + .filter(|value| !value.value().expose().is_empty()) .or(Some(key)) }), api_base: inputs.api_base.and_then(|base| { @@ -592,12 +592,17 @@ impl AzureDocumentIntelligenceOcrConfig { )?; return Ok(connection.extra_headers.clone()); } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(self.get_api_key_env_var().and_then(env_lookup)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); + let key = nonblank( + connection + .api_key + .as_ref() + .map(|key| key.expose().to_string()), + ) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); if let Some(key) = key { super::super::common_utils::validate_destination(connection, key.source())?; return Ok( @@ -796,7 +801,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 2012f740173..7ef051e8986 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -142,12 +142,17 @@ impl AzureAiOcrConfig { super::common_utils::validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(self.get_api_key_env_var().and_then(env_lookup)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); + let key = nonblank( + connection + .api_key + .as_ref() + .map(|key| key.expose().to_string()), + ) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); if let Some(key) = key { super::common_utils::validate_destination(connection, key.source())?; return Ok(bearer_headers(connection, key.value())); @@ -196,7 +201,7 @@ mod tests { #[fixture] fn connection() -> OcrConnection { OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_base: Some("https://example.com".into()), ..Default::default() } @@ -288,7 +293,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index e6fe5d9556d..8321dcfb4ce 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, future::Future, time::Duration}; -use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; +use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; use litellm_core_utils::{ call_arguments::CallArguments, serde_compat::{FiniteF64, LaxI64}, @@ -90,21 +90,22 @@ pub enum OcrResponseFormat { #[derive(Clone, Default)] pub struct OcrCredentialInputs { - pub api_key: Option>, - pub dynamic_api_key: Option>, + pub api_key: Option>, + pub dynamic_api_key: Option>, pub api_base: Option>, pub dynamic_api_base: Option>, } impl OcrCredentialInputs { pub fn new( - api_key: Option, + api_key: Option, api_key_source: InputSource, api_base: Option, api_base_source: InputSource, ) -> Self { Self { - api_key: nonblank(api_key).map(|value| Sourced::new(value, api_key_source)), + api_key: nonblank(api_key.as_ref().map(|key| key.expose().to_string())) + .map(|value| Sourced::new(SecretValue::new(value), api_key_source)), dynamic_api_key: None, api_base: nonblank(api_base).map(|value| Sourced::new(value, api_base_source)), dynamic_api_base: None, @@ -159,7 +160,7 @@ fn nonblank(value: Option) -> Option { #[derive(Clone)] pub struct OcrConnection { - pub api_key: Option, + pub api_key: Option, pub api_key_source: InputSource, pub api_base: Option, pub api_base_source: InputSource, @@ -209,7 +210,7 @@ impl Default for OcrConnection { #[derive(Clone, Default)] pub struct ResolvedOcrCredentials { - pub api_key: Option>, + pub api_key: Option>, pub api_base: Option>, } @@ -428,7 +429,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { ResolvedOcrCredentials { api_key: inputs .dynamic_api_key - .filter(|value| !value.value().is_empty()) + .filter(|value| !value.value().expose().is_empty()) .or(inputs.api_key), api_base: inputs .dynamic_api_base diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index f353c22d8c4..2528c967f41 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -179,8 +179,8 @@ impl CohereParseConfig { } let key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -718,7 +718,7 @@ mod tests { assert!(matches!( CohereParseConfig.resolve_headers( &OcrConnection { - api_key: Some(" ".into()), + api_key: Some(litellm_auth::SecretValue::new(" ")), ..Default::default() }, &|_| None, diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index e93ddee3c50..7635bdd3d04 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -3,9 +3,9 @@ use std::{sync::OnceLock, time::Duration}; use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; -use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_callbacks::event::WireRequest; use serde::{Serialize, de::DeserializeOwned}; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::{ base_llm::ocr::{ @@ -26,11 +26,7 @@ use crate::{ /// The route's view of one call, handed to provider code that has to reach the /// caller's hooks mid-flight (guardrails on the outgoing body, raw response events). pub trait CallHooks: Send + Sync { - fn before_send( - &self, - wire: WireRequest, - passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result>; + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result>; fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), E>>; } @@ -231,9 +227,8 @@ pub async fn transform_request_body( config.get_supported_ocr_params(&request.model), )?; config.validate_request_body(&composed)?; - let passthrough_fields = Passthrough::unchanged(&caller_inputs(request)?, &composed); let changed = hooks - .before_send(wire_request(url, headers, composed), passthrough_fields) + .before_send(wire_request(url, headers, composed)) .await?; if !changed.body.is_object() { return Err(Error::RequestField { @@ -252,21 +247,6 @@ fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireReq } } -fn caller_inputs(request: &PreparedOcrRequest) -> Result, Error> { - let document = request - .caller_document - .then(|| serde_json::to_value(&request.document)) - .transpose() - .map_err(|_| Error::RequestField { - path: "document".into(), - })?; - let params: Map = request.optional_params.clone().into(); - Ok(params - .into_iter() - .chain(document.map(|document| ("document".to_string(), document))) - .collect()) -} - pub fn build_http_request( client: &OcrClient, request: &PreparedOcrRequest, @@ -294,9 +274,7 @@ pub async fn guardrail_document( let body = serde_json::to_value(&request.document).map_err(|_| Error::RequestField { path: "document".into(), })?; - let changed = hooks - .before_send(wire_request(url, headers, body), Passthrough::default()) - .await?; + let changed = hooks.before_send(wire_request(url, headers, body)).await?; let document = decode_request_value(changed.body, "guardrail.document")?; Ok((document, changed.headers)) } diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index c2038d0552d..9028f09c5ab 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -135,8 +135,8 @@ impl MistralOcrConfig { } let api_key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -212,7 +212,7 @@ mod tests { #[default(vec![])] extra_headers: Vec<(String, String)>, ) -> OcrConnection { OcrConnection { - api_key: api_key.map(str::to_string), + api_key: api_key.map(litellm_auth::SecretValue::new), extra_headers, ..OcrConnection::default() } diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index ca2bae9c3bb..ec876fafb8f 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -442,8 +442,8 @@ fn resolve_headers( } let api_key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -629,7 +629,7 @@ mod tests { #[test] fn explicit_key_precedes_environment_key() { let connection = OcrConnection { - api_key: Some("passed-key".into()), + api_key: Some(litellm_auth::SecretValue::new("passed-key")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); @@ -639,7 +639,7 @@ mod tests { #[test] fn blank_explicit_key_uses_environment_key() { let connection = OcrConnection { - api_key: Some(" ".into()), + api_key: Some(litellm_auth::SecretValue::new(" ")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index ea0bcf3d08c..c2cb23d0010 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -134,7 +134,10 @@ impl VertexAiOcrConfig { .vertex_auth() .validate_environment( connection.extra_headers.clone(), - connection.api_key.as_deref(), + connection + .api_key + .as_ref() + .map(litellm_auth::SecretValue::expose), config, &credential_env, ) diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 9932594e2f5..5dccfb4aca8 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,7 +1,7 @@ - Target invariants, not completion claims; these supersede older conflicting bridge guidance - Keep this crate the product-specific PyO3 consumer of `litellm-host-python` - Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract - - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, `passthrough_fields` re-aliasing) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy + - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy - Value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment live in `litellm-host-python`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy` owns `Logging` dispatch policy - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 9dc891a91d6..77212e3d38e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,9 +1,9 @@ use litellm_auth::ResolvedCredential; use litellm_core::ocr::route::{Ocr, OcrOp, OcrOpResult}; -use litellm_host_python::{RouteHost, missing_state, to_py}; +use litellm_host_python::{HostOpError, RouteHost, missing_state, to_py}; use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}; use pyo3::{ - exceptions::PyBaseException, + exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, types::PyDict, @@ -57,12 +57,8 @@ impl OcrRouteHost { .ok_or_else(missing_state)? .acquire(py) } -} -impl RouteHost for OcrRouteHost { - type Route = Ocr; - - fn invoke( + fn answer( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, @@ -88,6 +84,40 @@ impl RouteHost for OcrRouteHost { } } + fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { + if !error.is_instance_of::(py) { + return error; + } + let provider = match &self.data { + OcrHostData::Projected(handles) => handles.provider, + _ => "", + }; + let mapped = py + .import("litellm.rust_bridge.ocr.route_host") + .and_then(|module| module.getattr("map_failure")) + .and_then(|map| map.call1((error.value(py), self.request.bind(py), provider))) + .and_then(|mapped| mapped.extract::>().map_err(PyErr::from)); + match mapped { + Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), + Err(_) => error, + } + } +} + +impl RouteHost for OcrRouteHost { + type Route = Ocr; + type Failure = PyErr; + + fn invoke( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + op: OcrOp, + ) -> Result> { + self.answer(py, arguments, op) + .map_err(|error| HostOpError::Python(self.map_failure(py, error))) + } + fn complete(&mut self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult> { py.import("litellm.rust_bridge.ocr.route_host")? .getattr("response")? @@ -95,27 +125,14 @@ impl RouteHost for OcrRouteHost { .map(Bound::unbind) } - fn native_error(error: Error) -> PyErr { - ocr_error_to_pyerr(error) + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + Ok(self.map_failure(py, ocr_error_to_pyerr(error))) } fn host_error(error: &PyErr) -> Error { Error::InvalidRequest(error.to_string()) } - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult { - let provider = match &self.data { - OcrHostData::Projected(handles) => handles.provider, - _ => "", - }; - let mapped: Py = py - .import("litellm.rust_bridge.ocr.route_host")? - .getattr("map_failure")? - .call1((error.value(py), self.request.bind(py), provider))? - .extract()?; - Ok(PyErr::from_value(mapped.into_bound(py).into_any())) - } - fn close(&mut self, _: Python<'_>) { self.data = OcrHostData::Released; } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 7ffa129f85c..5dd2aa804b8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,3 +1,4 @@ +use litellm_auth::SecretValue; use litellm_core::ocr::{ types::{LiteLLMOcrRequest, OcrDocumentInput}, wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, @@ -31,7 +32,7 @@ struct OcrArguments<'a, 'py> { impl<'py> OcrArguments<'_, 'py> { fn lookup(&self, name: &str) -> PyResult> { - litellm_callbacks_legacy::lookup(self.kwargs, self.request, name)? + litellm_host_python::lookup(self.kwargs, self.request, name)? .ok_or_else(|| PyValueError::new_err(format!("missing argument: {name}"))) } @@ -47,8 +48,11 @@ impl<'py> OcrArguments<'_, 'py> { self.lookup("document") } - fn api_key(&self) -> PyResult> { - self.lookup("api_key")?.extract() + fn api_key(&self) -> PyResult> { + Ok(self + .lookup("api_key")? + .extract::>()? + .map(SecretValue::new)) } fn api_base(&self) -> PyResult> { diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4a9a65b1485..7248c2f3590 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1265,11 +1265,6 @@ class Logging(LiteLLMLoggingBaseClass): additional_args.get("api_base", "") ) - def record_api_call_start_time(self) -> None: - self.model_call_details["api_call_start_time"] = datetime.datetime.now() - if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] - def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API try: @@ -1334,7 +1329,15 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - self.record_api_call_start_time() + self.model_call_details["api_call_start_time"] = datetime.datetime.now() + # Set-once first provider-handoff instant. api_call_start_time + # is overwritten on every retry, so it can't measure one-time + # preprocessing; pinning the first attempt excludes retry loops + # + backoff. Logging object only — must NOT go into + # litellm_params["metadata"] (caller request metadata, typed + # Dict[str, str], echoed downstream; a datetime breaks it). + if self.model_call_details.get("first_api_call_start_time") is None: + self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) for callback in callbacks: @@ -1468,21 +1471,16 @@ class Logging(LiteLLMLoggingBaseClass): """ return _get_masked_values(headers, ignore_sensitive_values=ignore_sensitive_headers) - def record_post_call( - self, original_response: object, input: object, api_key: object, additional_args: dict[str, object] - ) -> None: - self.model_call_details["input"] = input - self.model_call_details["api_key"] = api_key - self.model_call_details["original_response"] = original_response - self.model_call_details["additional_args"] = additional_args - self.model_call_details["log_event_type"] = "post_api_call" - def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received if isinstance(original_response, dict): original_response = json.dumps(original_response, default=str) try: - self.record_post_call(original_response, input, api_key, additional_args) + self.model_call_details["input"] = input + self.model_call_details["api_key"] = api_key + self.model_call_details["original_response"] = original_response + self.model_call_details["additional_args"] = additional_args + self.model_call_details["log_event_type"] = "post_api_call" attr: Literal["warning", "debug"] if self.litellm_request_debug: @@ -2177,7 +2175,6 @@ class Logging(LiteLLMLoggingBaseClass): logging_result, start_time, end_time, - build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params: Final = getattr(logging_result, "_hidden_params", {}) @@ -2202,9 +2199,6 @@ class Logging(LiteLLMLoggingBaseClass): else: self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result) - if not build_logging_payload: - return - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( logging_result, start_time, end_time ) @@ -2266,7 +2260,6 @@ class Logging(LiteLLMLoggingBaseClass): end_time=None, cache_hit=None, standard_logging_object: StandardLoggingPayload | None = None, - build_logging_payload: bool = True, ): try: if start_time is None: @@ -2304,7 +2297,6 @@ class Logging(LiteLLMLoggingBaseClass): logging_result=logging_result, start_time=start_time, end_time=end_time, - build_logging_payload=build_logging_payload, ) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = standard_logging_object @@ -3328,9 +3320,7 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) - def _failure_handler_helper_fn( - self, exception, traceback_exception, start_time=None, end_time=None, build_logging_payload: bool = True - ): + def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: start_time = self.start_time if end_time is None: @@ -3365,9 +3355,6 @@ class Logging(LiteLLMLoggingBaseClass): metadata: Final = self.model_call_details["litellm_params"].get("metadata", {}) or {} metadata.update(exception.headers) - if not build_logging_payload: - return start_time, end_time - ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = get_standard_logging_object_payload( diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index e05d9368fa8..406dc55cfea 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -6,24 +6,22 @@ registries it fans out to. It expires with that contract. from __future__ import annotations +import contextvars import datetime -import os +import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Coroutine, Mapping from dataclasses import dataclass from typing import ( TYPE_CHECKING, Final, - Literal, Protocol, - TypeAlias, cast, # noqa: TID251 # bounded compatibility calls into legacy Python integrations ) -from typing_extensions import assert_never - if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CredentialItem class MetadataUpdater(Protocol): @@ -42,7 +40,6 @@ class MetadataUpdater(Protocol): class CallSetup: logger: Logging kwargs: dict[str, object] - bridge_owned: bool def setup( @@ -61,9 +58,9 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return CallSetup(supplied, arguments, bridge_owned=False) + return CallSetup(supplied, arguments) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return CallSetup(logger, prepared, bridge_owned=True) + return CallSetup(logger, prepared) def check_limits(kwargs: Mapping[str, object]) -> None: @@ -93,87 +90,219 @@ def finalize( update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) -def deployment_callbacks_needed() -> bool: - import litellm - from litellm.integrations.custom_logger import CustomLogger +class LoggingSurface(Protocol): + def update_from_kwargs( + self, + kwargs: dict[str, object], + litellm_params: dict[str, object] | None = None, + optional_params: dict[str, object] | None = None, + model: str | None = None, + user: str | None = None, + **additional_params: object, + ) -> None: ... - return any(isinstance(callback, CustomLogger) for callback in litellm.callbacks) + def pre_call( + self, input: object, api_key: object, model: object = None, additional_args: dict[str, object] = ... + ) -> object: ... + + def post_call( + self, + original_response: object, + input: object = None, + api_key: object = None, + additional_args: dict[str, object] = ..., + ) -> object: ... + + def handle_sync_success_callbacks_for_async_calls( + self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: object = None + ) -> None: ... + + def failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: ... + + def async_failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> Coroutine[object, object, None]: ... + + def success_handler( + self, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: object, + ) -> None: ... + + def async_success_handler( + self, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: object, + ) -> Coroutine[object, object, None]: ... -Phase: TypeAlias = Literal[ - "input", "sync_success", "sync_success_async", "async_success", "sync_failure", "async_failure", "payload" -] +if TYPE_CHECKING: + _LOGGING_CONFORMS: type[LoggingSurface] = Logging -def callbacks_needed(logger: Logging, phase: Phase) -> bool: - import litellm - from litellm._logging import ( - _is_debugging_on, # pyright: ignore[reportPrivateUsage] # use the same debug gate as Logging +class LoggingWorker(Protocol): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... + + +class DeploymentHook(Protocol): + def __call__(self, kwargs: dict[str, object], call_type: str) -> Awaitable[object]: ... + + +class DeploymentSuccessHook(Protocol): + def __call__(self, request_data: dict[str, object], response: object, call_type: object) -> Awaitable[object]: ... + + +class DeploymentFailureHook(Protocol): + def __call__(self, request_data: Mapping[str, object], exception: Exception, call_type: str) -> Awaitable[None]: ... + + +def update_logging( + logger: LoggingSurface, + kwargs: dict[str, object], + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + custom_llm_provider: str, +) -> None: + logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, ) - if ( - _is_debugging_on() - or getattr(logger, "litellm_request_debug", False) - or os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD") - ): - return True - input_needed: Final = bool( - litellm.input_callback - or litellm._async_input_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_input_callbacks - or callable(getattr(logger, "logger_fn", None)) - or logger.log_raw_request_response - or litellm.log_raw_request_response + +def pre_call(logger: LoggingSurface, input: str, api_key: str | None, additional_args: dict[str, object]) -> None: + logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) + + +def post_call( + logger: LoggingSurface, original_response: str, api_key: str | None, additional_args: dict[str, object] +) -> None: + logger.post_call(original_response=original_response, api_key=api_key, additional_args=additional_args) + + +def defers_async_logging(logger: LoggingSurface) -> bool: + return bool(getattr(logger, "_defer_async_logging", False)) + + +def defer_success(logger: LoggingSurface, pending: object) -> None: + setattr(logger, "_native_pending_logging", pending) + + +def sync_success_for_async_call( + logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime +) -> None: + logger.handle_sync_success_callbacks_for_async_calls(result=response, start_time=start, end_time=end) + + +def failure_handler( + logger: LoggingSurface, error: Exception, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +) -> Coroutine[object, object, None] | None: + trace: Final = "".join(traceback.format_exception(error)) + if asynchronous: + return logger.async_failure_handler(error, trace, start, end) + logger.failure_handler(error, trace, start, end) + return None + + +def submit_success(logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime) -> None: + from litellm.litellm_core_utils.litellm_logging import executor + + executor.submit(contextvars.copy_context().run, logger.success_handler, response, start, end) + + +def async_success_handler( + logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime +) -> Coroutine[object, object, None]: + return logger.async_success_handler(response, start, end) + + +def enqueue_logging(coroutine: Coroutine[object, object, None]) -> None: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + worker: Final = cast( # cast-ok: bounded adapter for the untyped logging worker + LoggingWorker, GLOBAL_LOGGING_WORKER ) - match phase: - case "input": - return input_needed - case "sync_success": - return bool(litellm.success_callback or logger.dynamic_success_callbacks) - case "sync_success_async": - return bool( - (litellm.success_callback or logger.dynamic_success_callbacks) - and logger._should_run_sync_callbacks_for_async_calls() # pyright: ignore[reportPrivateUsage] # preserve async call filtering of sync callbacks - ) - case "async_success": - return bool(litellm._async_success_callback or logger.dynamic_async_success_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "sync_failure": - return bool(litellm.failure_callback or logger.dynamic_failure_callbacks) - case "async_failure": - return bool(litellm._async_failure_callback or logger.dynamic_async_failure_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "payload": - return bool( - input_needed - or litellm.success_callback - or litellm.failure_callback - or litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or litellm._async_failure_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_success_callbacks - or logger.dynamic_async_success_callbacks - or logger.dynamic_failure_callbacks - or logger.dynamic_async_failure_callbacks - ) - case _: - assert_never(phase) + contextvars.copy_context().run(worker.ensure_initialized_and_enqueue, coroutine) -def success_bookkeeping( - logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_success" if asynchronous else "sync_success" - if logger.should_run_logging(phase): - logger._success_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain success bookkeeping without constructing a callback payload - result=response, start_time=start, end_time=end, build_logging_payload=False - ) - logger.has_run_logging(phase) +def restore_context(logger: LoggingSurface) -> None: + from litellm.utils import ( + _restore_correlation_context_if_supported, # pyright: ignore[reportPrivateUsage] # the @client wrapper restores the same correlation context + ) + + _restore_correlation_context_if_supported(logger) -def failure_bookkeeping( - logger: Logging, error: BaseException, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_failure" if asynchronous else "sync_failure" - if logger.should_run_logging(phase): - logger._failure_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain failure accounting without formatting an unused traceback or payload - error, "", start, end, build_logging_payload=False - ) - logger.has_run_logging(phase) +def custom_pricing_fields() -> tuple[str, ...]: + from litellm.types.utils import CustomPricingLiteLLMParams + + return tuple(CustomPricingLiteLLMParams.model_fields) + + +def is_internal_call() -> bool: + from litellm._internal_context import is_internal_call as internal + + return internal.get() + + +def credential_list() -> list[CredentialItem]: + import litellm + + return litellm.credential_list + + +def warn_unknown_credential(name: str, loaded: int) -> None: + from litellm._logging import verbose_logger + + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + name, + loaded, + ) + + +def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitable[object]: + from litellm import utils + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentHook, utils.async_pre_call_deployment_hook + ) + return hook(kwargs, call_type) + + +def after_deployment_success(kwargs: dict[str, object], response: object, call_type: str) -> Awaitable[object]: + from litellm import utils + from litellm.types.utils import CallTypes + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentSuccessHook, utils.async_post_call_success_deployment_hook + ) + return hook(kwargs, response, CallTypes(call_type)) + + +def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_type: str) -> Awaitable[None]: + from litellm import utils + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentFailureHook, utils.async_post_call_failure_deployment_hook + ) + return hook(kwargs, error, call_type) diff --git a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py b/tests/test_litellm/rust_bridge/test_legacy_callbacks.py index a4474c85230..a0906c7c5be 100644 --- a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py +++ b/tests/test_litellm/rust_bridge/test_legacy_callbacks.py @@ -1,12 +1,16 @@ import datetime +import inspect from collections.abc import Mapping +from pathlib import Path from types import MappingProxyType from typing import Final import pytest +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.rust_bridge import legacy_callbacks as legacy from litellm.rust_bridge.legacy_callbacks import check_limits, setup _OCR_KWARGS: Final = MappingProxyType( @@ -56,13 +60,12 @@ def _supplied_logger() -> Logging: ) -def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None: +def test_setup_reuses_a_supplied_logger() -> None: supplied: Final = _supplied_logger() result: Final = setup( "aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True ) assert result.logger is supplied - assert result.bridge_owned is False @pytest.mark.parametrize( @@ -73,7 +76,15 @@ def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None: ], ids=["ocr", "embedding"], ) -def test_setup_owns_every_logger_it_builds(call_type: str, kwargs: Mapping[str, object]) -> None: +def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Mapping[str, object]) -> None: result: Final = setup(call_type, (), kwargs, datetime.datetime.now(), asynchronous=True) - assert result.bridge_owned is True assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"] + + +CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy/python_contract.json" + + +def test_the_rust_contract_matches_the_shim_signatures() -> None: + contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) + + assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 085ea4a14c0..5fca927bea3 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -4,7 +4,7 @@ import gc import json import threading import weakref -from collections.abc import Coroutine +from collections.abc import Awaitable, Callable, Coroutine from contextvars import ContextVar from typing import Final @@ -400,7 +400,7 @@ async def test_response_limit_is_enforced_at_the_public_boundary(ocr_server: Rec @pytest.mark.asyncio @pytest.mark.parametrize("failure", [False, True]) -async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( +async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_client_wrapper( ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, failure: bool, @@ -414,8 +414,12 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( submissions = 0 enqueues = 0 - def deployment(self, *args: object, **kwargs: object) -> None: - self.deployments += 1 + def counting(self, hook: Callable[..., Awaitable[object]]) -> Callable[..., Awaitable[object]]: + async def counted(*args: object, **kwargs: object) -> object: + self.deployments += 1 + return await hook(*args, **kwargs) + + return counted def submit(self, *args: object, **kwargs: object) -> None: self.submissions += 1 @@ -430,7 +434,7 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( "async_post_call_success_deployment_hook", "async_post_call_failure_deployment_hook", ): - monkeypatch.setattr(utils, name, probe.deployment) + monkeypatch.setattr(utils, name, probe.counting(getattr(utils, name))) monkeypatch.setattr(litellm_logging, "executor", probe) monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) if failure: @@ -447,17 +451,16 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( assert response._hidden_params["response_cost"] is not None assert response._hidden_params["_response_ms"] > 0 assert trace_id_var.get() == "callback-free-parent" - assert probe.deployments == probe.submissions == probe.enqueues == 0 + assert probe.deployments == 2 + assert probe.submissions == probe.enqueues == 0 assert len(created_loggers) == 1 logger: Final = created_loggers[0] - assert not hasattr(logger, "_native_pending_logging") - assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] - assert "standard_logging_object" not in logger.model_call_details - assert ( - "original_response" not in logger.model_call_details or logger.model_call_details["original_response"] is None - ) - assert "complete_input_dict" not in logger.model_call_details.get("additional_args", {}) - assert logger.model_call_details["response_cost"] == (0 if failure else response._hidden_params["response_cost"]) + if failure: + assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] + assert logger.model_call_details["response_cost"] == 0 + else: + assert getattr(logger, "_native_pending_logging", None) is not None + assert "end_time" not in logger.model_call_details @pytest.mark.asyncio From a737e3430a979b78c4d84ceac6d7605aca729d79 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 14:41:50 -0700 Subject: [PATCH 088/102] test(rust): property and parametrized tests for legacy callback contracts The payload boundary of callbacks-legacy gets a model-based proptest: for any JSON body, caller keywords and callback edit, keywords the route sends unchanged reach pre_call as the caller's own objects, and the wire is the body pre_call received as the callback left it. A parametrized test pins that a keyword the bridge never reads keeps its identity through setup, the deployment hook, check_limits and prepare. Behaviour owned by the real Logging object is pinned end to end in the OCR tests: a hypothesis version of the body property over HTTP, sync hooks seeing no running event loop, retained payloads staying intact after the call, success callbacks sharing one standard logging payload, and state stashed before a blocking deployment hook raises reaching both failure callback families --- litellm-rust/Cargo.toml | 1 + .../crates/callbacks-legacy/Cargo.toml | 1 + .../tests/deployment_hooks.rs | 37 ++++ .../crates/callbacks-legacy/tests/payload.rs | 175 +++++++++++++++++- tests/test_litellm_rust/ocr/test_callbacks.py | 163 +++++++++++++++- 5 files changed, 370 insertions(+), 7 deletions(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index ffdbf64bb49..32f925b8d8b 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -26,6 +26,7 @@ litellm-token-counter = { path = "crates/token-counter" } litellm-host-python = { path = "crates/host-python" } bytes = "1" +proptest = "1.7.0" pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy/Cargo.toml index 3cf9382f857..efacde051ed 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy/Cargo.toml @@ -16,5 +16,6 @@ serde_json.workspace = true [dev-dependencies] litellm-auth.workspace = true +proptest.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index ea3510de17e..7bad09c7890 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -97,6 +97,43 @@ assert checked is prepared }); } +#[rstest] +#[case::synchronous(false)] +#[case::asynchronous(true)] +fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( + #[case] asynchronous: bool, +) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +opaque = object() +hooked = [] +logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} +kwargs = {'logger': logger, 'vendor_extension': opaque} +", + ); + let (mut logging, step) = begin(py, &locals, asynchronous); + let step = match step { + LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), + step => step, + }; + locals.set_item("prepared", arguments(py, step)).unwrap(); + locals.set_item("asynchronous", asynchronous).unwrap(); + run( + py, + &locals, + c" +assert prepared['vendor_extension'] is opaque +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked['vendor_extension'] is opaque +assert hooked == ([opaque] if asynchronous else []), hooked +", + ); + }); +} + #[test] fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { Python::initialize(); diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 68c0a2b1e15..5c49383bf37 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -2,10 +2,11 @@ use std::ffi::CStr; use litellm_auth::SecretValue; use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{LifecycleStep, PythonLifecycle}; +use litellm_host_python::{LifecycleStep, PythonLifecycle, to_py}; +use proptest::prelude::*; use pyo3::prelude::*; use rstest::rstest; -use serde_json::{Value, json}; +use serde_json::{Map, Value, json}; use super::LegacyLogging; use crate::PythonLogger; @@ -57,10 +58,24 @@ fn before_send_with_secrets( optional_params: Value, body: Value, secret_fields: &[&str], +) -> WireRequest { + before_send_bound(&[], script, optional_params, body, secret_fields) +} + +/// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. +fn before_send_bound( + bindings: &[(&str, &Value)], + script: &CStr, + optional_params: Value, + body: Value, + secret_fields: &[&str], ) -> WireRequest { Python::initialize(); Python::attach(|py| { let locals = namespace(py, PAYLOAD_LOGGER); + for &(name, value) in bindings { + locals.set_item(name, to_py(py, value).unwrap()).unwrap(); + } run(py, &locals, script); let mut logging = LegacyLogging { logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), @@ -346,3 +361,159 @@ def check(): json!({"document": document(DOCUMENT), "include_image_base64": true}) ); } + +/// What one pre-call callback does to the payload it is handed. +#[derive(Clone, Debug)] +enum Edit { + Nothing, + Set(String, Value), + Remove(String), + Rebind(Value), + RebindThenSetRetained(String, Value), +} + +impl Edit { + fn script(&self) -> Value { + match self { + Self::Nothing => json!({"kind": "nothing"}), + Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), + Self::Remove(key) => json!({"kind": "remove", "key": key}), + Self::Rebind(value) => json!({"kind": "rebind", "value": value}), + Self::RebindThenSetRetained(key, value) => { + json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) + } + } + } + + /// The legacy contract: the provider is sent the body object `pre_call` received, as + /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and + /// leaves that object alone. + fn sent(&self, body: &Map) -> Value { + let mut sent = body.clone(); + match self { + Self::Nothing | Self::Rebind(_) => {} + Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { + sent.insert(key.clone(), value.clone()); + } + Self::Remove(key) => { + sent.remove(key); + } + } + Value::Object(sent) + } +} + +/// How the caller's keyword for a body key relates to what the route sends under it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Caller { + PassedUnchanged, + RewrittenByTheRoute, + NotPassed, +} + +const MODEL: &CStr = c" +aliased = {} +def on_pre_call(args): + body = args['complete_input_dict'] + aliased.update({name: body[name] is kwargs[name] for name in unchanged}) + kind = edit['kind'] + if kind == 'set': + body[edit['key']] = edit['value'] + elif kind == 'remove': + body.pop(edit['key'], None) + elif kind == 'rebind': + args['complete_input_dict'] = edit['value'] + elif kind == 'rebind_then_set_retained': + args['complete_input_dict'] = {} + body[edit['key']] = edit['value'] +def check(): + assert aliased == {name: True for name in unchanged}, aliased + assert logger.names() == ['pre_call', 'post_call'], logger.calls +"; + +fn json_value() -> impl Strategy { + let leaf = prop_oneof![ + Just(Value::Null), + any::().prop_map(Value::from), + any::().prop_map(Value::from), + any::() + .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) + .prop_map(Value::from), + ".{0,8}".prop_map(Value::from), + ]; + leaf.prop_recursive(3, 24, 4, |inner| { + prop_oneof![ + prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), + prop::collection::btree_map(key(), inner, 0..4) + .prop_map(|fields| Value::Object(fields.into_iter().collect())), + ] + }) +} + +fn key() -> impl Strategy { + "[a-z]{1,6}" +} + +fn caller() -> impl Strategy { + prop_oneof![ + Just(Caller::PassedUnchanged), + Just(Caller::RewrittenByTheRoute), + Just(Caller::NotPassed), + ] +} + +fn edit() -> impl Strategy { + prop_oneof![ + Just(Edit::Nothing), + (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), + key().prop_map(Edit::Remove), + json_value().prop_map(Edit::Rebind), + (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), + ] +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(128))] + + /// For any body, any caller keywords and any callback edit: every keyword the route + /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is + /// sent exactly what the model says, so a callback that edits nothing changes nothing. + #[test] + fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( + fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), + edit in edit(), + ) { + let body: Map = fields + .iter() + .map(|(name, (value, _))| (name.clone(), value.clone())) + .collect(); + let kwargs: Map = fields + .iter() + .filter_map(|(name, (value, caller))| match caller { + Caller::PassedUnchanged => Some((name.clone(), value.clone())), + Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), + Caller::NotPassed => None, + }) + .collect(); + let unchanged: Value = fields + .iter() + .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) + .map(|(name, _)| Value::from(name.clone())) + .collect(); + + let wire = before_send_bound( + &[ + ("kwargs", &Value::Object(kwargs)), + ("unchanged", &unchanged), + ("edit", &edit.script()), + ], + MODEL, + json!({}), + Value::Object(body.clone()), + &[], + ); + + prop_assert_eq!(wire.body, edit.sent(&body)); + prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } +} diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 27cdcc4d997..7cc19b8c090 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -1,24 +1,28 @@ import asyncio import copy +import gc import queue import threading from typing import Final import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, OCR_RESPONSE, + call_native, call_native_aocr, call_native_ocr, request_body, request_headers, ) -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension @@ -123,9 +127,7 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_ "callbacks": [Retain(), Edit()], } response: Final = ( - await call_native_aocr(ocr_server, **arguments) - if asynchronous - else call_native_ocr(ocr_server, **arguments) + await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) ) assert aliases == [True] @@ -291,6 +293,157 @@ def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registere assert "log_failure_event" not in recorder.names +JSON_SCALARS: Final = ( + st.none() + | st.booleans() + | st.integers(min_value=-(2**63), max_value=2**63 - 1) + | st.floats(allow_nan=False, allow_infinity=False) + | st.text(max_size=8) +) +JSON_VALUES: Final = st.recursive( + JSON_SCALARS, + lambda children: st.lists(children, max_size=3) | st.dictionaries(st.text(max_size=6), children, max_size=3), + max_leaves=8, +) + + +LATEST_EDITS: Final[list[dict[str, object]]] = [] + + +class ApplyLatestEdits(CustomLogger): + """Registrations can outlive one hypothesis example, so every instance applies the current example's edits.""" + + def __init__(self, latest: list[dict[str, object]]) -> None: + super().__init__() + self.latest = latest + + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs).update(copy.deepcopy(self.latest[-1])) + + +@settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) +@given(edits=st.dictionaries(st.from_regex(r"x_[a-z]{1,6}", fullmatch=True), JSON_VALUES, max_size=3)) +def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it( + ocr_server: RecordingServer, edits: dict[str, object] +) -> None: + LATEST_EDITS.append(edits) + + call_native_ocr_with_callbacks(ocr_server, [ApplyLatestEdits(LATEST_EDITS)]) + + assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits} + + +@pytest.mark.parametrize("hook", ["log_pre_api_call", "logging_hook", "log_success_event"]) +def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingServer, hook: str) -> None: + recorder: Final = RecordingLogger() + + call_native_ocr_with_callbacks(ocr_server, [recorder]) + + [event] = recorder.wait_for(hook) + assert event.loop is None + assert (event.thread is threading.current_thread()) == (hook == "log_pre_api_call") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + retained: Final = [] + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + retained.append((kwargs, request_body(kwargs), request_headers(kwargs))) + + await call_native(ocr_server, asynchronous, callbacks=[Retain()]) + await drain_logging() + gc.collect() + + [(details, body, headers)] = retained + assert body == ocr_server.requests[0].body + assert headers + assert all(ocr_server.requests[0].headers[name] == value for name, value in headers.items()) + assert details["additional_args"]["complete_input_dict"] is body + assert details["additional_args"]["headers"] is headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("family", ["sync", "async"]) +async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server: RecordingServer, family: str) -> None: + queued: Final = [] + finished: Final = threading.Event() + + def queue_payload(kwargs: dict[str, object]) -> None: + queued.append(kwargs["standard_logging_object"]) + + def strip_payload(kwargs: dict[str, object]) -> None: + payload: Final = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + payload["stripped-by-a-later-callback"] = True + finished.set() + + class QueuePayload(CustomLogger): + if family == "sync": + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + queue_payload(kwargs) + + else: + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + queue_payload(kwargs) + + class StripPayload(CustomLogger): + if family == "sync": + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + strip_payload(kwargs) + + else: + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + strip_payload(kwargs) + + await call_native(ocr_server, family == "async", callbacks=[QueuePayload(), StripPayload()]) + await drain_logging() + + assert await asyncio.to_thread(finished.wait, 10) + assert [payload["stripped-by-a-later-callback"] for payload in queued] == [True] + + +@pytest.mark.asyncio +async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks( + ocr_server: RecordingServer, +) -> None: + token: Final = object() + observed: Final = [] + + class Blocked(Exception): + pass + + class Block(CustomLogger): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + request_data["litellm_logging_obj"].model_call_details["blocked-by"] = token + raise Blocked("blocked after the provider answered") + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("success", None, None)) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs.get("blocked-by"), kwargs["exception"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs.get("blocked-by"), kwargs["exception"])) + + litellm.callbacks.append(Block()) + + with pytest.raises(Blocked) as raised: + await call_native_aocr(ocr_server) + await drain_logging() + + assert observed == [("sync", token, raised.value), ("async", token, raised.value)] + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context( From 30f7b8442bd8b7258ad4bc5a41a9b85546d7a83c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 14:42:59 -0700 Subject: [PATCH 089/102] chore(rust): lock proptest --- litellm-rust/Cargo.lock | 62 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 0265e0adbc2..bf58a81a6c5 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2044,6 +2044,7 @@ dependencies = [ "litellm-auth", "litellm-callbacks", "litellm-host-python", + "proptest", "pyo3", "rstest", "serde_json", @@ -2601,6 +2602,25 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "proptest" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" +dependencies = [ + "bit-set", + "bit-vec", + "bitflags", + "num-traits", + "rand 0.9.5", + "rand_chacha 0.9.0", + "rand_xorshift", + "regex-syntax", + "rusty-fork", + "tempfile", + "unarray", +] + [[package]] name = "pyo3" version = "0.29.2" @@ -2682,6 +2702,12 @@ dependencies = [ "serde", ] +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + [[package]] name = "quinn" version = "0.11.11" @@ -2845,6 +2871,15 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rand_xorshift" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" +dependencies = [ + "rand_core 0.9.5", +] + [[package]] name = "rayon" version = "1.12.0" @@ -3244,6 +3279,18 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" +[[package]] +name = "rusty-fork" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2" +dependencies = [ + "fnv", + "quick-error", + "tempfile", + "wait-timeout", +] + [[package]] name = "ryu" version = "1.0.23" @@ -4099,6 +4146,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "unarray" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + [[package]] name = "unicase" version = "2.9.0" @@ -4212,6 +4265,15 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "wait-timeout" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11" +dependencies = [ + "libc", +] + [[package]] name = "walkdir" version = "2.5.0" From ab27ee7efcbe982991095dbfbf5bd01662f8e69c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 14:44:19 -0700 Subject: [PATCH 090/102] test(rust): fix request count and header case in new OCR callback tests --- tests/test_litellm_rust/ocr/test_callbacks.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 7cc19b8c090..45b99d19d90 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -326,6 +326,7 @@ class ApplyLatestEdits(CustomLogger): def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it( ocr_server: RecordingServer, edits: dict[str, object] ) -> None: + ocr_server.expected_requests = None LATEST_EDITS.append(edits) call_native_ocr_with_callbacks(ocr_server, [ApplyLatestEdits(LATEST_EDITS)]) @@ -362,7 +363,7 @@ async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact( [(details, body, headers)] = retained assert body == ocr_server.requests[0].body assert headers - assert all(ocr_server.requests[0].headers[name] == value for name, value in headers.items()) + assert all(ocr_server.requests[0].headers[name.lower()] == value for name, value in headers.items()) assert details["additional_args"]["complete_input_dict"] is body assert details["additional_args"]["headers"] is headers From dd2e6c17bc19b839ff0a5e8500dfe40a7612f430 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 21:56:13 +0000 Subject: [PATCH 091/102] fix(rust): preserve OCR callback headers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/callbacks-legacy/src/adapter.rs | 12 ++++++++++-- .../crates/callbacks-legacy/src/callbacks.rs | 3 +++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 7204207cd62..53aa6ce9d2e 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -44,6 +44,7 @@ pub struct LegacyLogging { response: Option>, error: Option>, body: Option>, + headers: Option>, context: Option, asynchronous: bool, internal: bool, @@ -77,6 +78,7 @@ impl LegacyLogging { response: None, error: None, body: None, + headers: None, context: None, asynchronous, internal: false, @@ -245,6 +247,7 @@ impl PythonLifecycle for LegacyLogging { headers.set_item(name, value)?; } self.body = Some(body.clone().unbind()); + self.headers = Some(headers.clone().unbind()); self.context = Some(context.clone()); self.logger()?.pre_call( py, @@ -299,8 +302,13 @@ impl PythonLifecycle for LegacyLogging { .as_ref() .and_then(|context| context.api_key.as_ref()) .map(|api_key| api_key.expose()); - self.logger()? - .post_call(py, &raw.body, api_key, self.body.as_ref())?; + self.logger()?.post_call( + py, + &raw.body, + api_key, + self.body.as_ref(), + self.headers.as_ref(), + )?; Ok(LifecycleStep::Done) } (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs index 9464f1d6612..9fcfe98368e 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs @@ -38,6 +38,7 @@ pub trait LegacyCallbacks { original_response: &str, api_key: Option<&str>, body: Option<&Py>, + headers: Option<&Py>, ) -> PyResult<()>; fn defers_async_logging(&self, py: Python<'_>) -> bool; @@ -150,9 +151,11 @@ impl LegacyCallbacks for PythonLogger { original_response: &str, api_key: Option<&str>, body: Option<&Py>, + headers: Option<&Py>, ) -> PyResult<()> { let additional = PyDict::new(py); additional.set_item("complete_input_dict", body)?; + additional.set_item("headers", headers)?; Logging::PostCall.call( py, (self.object(py), original_response, api_key, &additional), From 4627ec4ea8f4ab6b53388f079bab54075d05a0ea Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 21:58:14 +0000 Subject: [PATCH 092/102] test(rust): cover post-call header identity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/callbacks-legacy/tests/payload.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 5c49383bf37..43128bc38ea 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -331,15 +331,19 @@ def on_pre_call(args): } #[test] -fn post_call_receives_the_raw_response_the_route_key_and_the_body_pre_call_saw() { +fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { before_send( c" def check(): original_response, api_key, additional_args = logger.post assert original_response == 'raw response', original_response assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) - assert additional_args == {'complete_input_dict': logger.pre['complete_input_dict']}, additional_args + assert additional_args == { + 'complete_input_dict': logger.pre['complete_input_dict'], + 'headers': logger.pre['headers'], + }, additional_args assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] + assert additional_args['headers'] is logger.pre['headers'] ", json!({"document": document(DOCUMENT)}), ); From 1a52bae77950ddab49f1d53a0a511fdaf1d7b570 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:32:28 -0700 Subject: [PATCH 093/102] add streaming message --- .../callbacks-legacy/python_contract.json | 17 ++ .../crates/callbacks-legacy/src/adapter.rs | 90 +++++++- .../callbacks-legacy/src/legacy_python.rs | 30 ++- .../crates/callbacks-legacy/tests/support.rs | 5 + litellm-rust/crates/callbacks/src/event.rs | 4 + litellm-rust/crates/callbacks/src/host.rs | 21 ++ litellm-rust/crates/callbacks/src/route.rs | 5 + litellm-rust/crates/callbacks/src/run.rs | 4 + litellm-rust/crates/core/src/machine/mod.rs | 17 +- .../crates/core/src/messages/handler.rs | 128 +++++------- litellm-rust/crates/core/src/messages/mod.rs | 36 +++- .../crates/core/src/messages/prepare.rs | 11 +- .../crates/core/src/messages/route.rs | 197 ++++++++++++++++++ litellm-rust/crates/core/src/ocr/route.rs | 2 + litellm-rust/crates/core/tests/ocr.rs | 2 + .../crates/host-python/src/adapter.rs | 8 + litellm-rust/crates/host-python/src/driver.rs | 77 ++++++- litellm-rust/crates/host-python/src/handle.rs | 13 ++ .../crates/python-bridge/src/marshal.rs | 16 -- .../python-bridge/src/routes/messages.rs | 88 -------- .../python-bridge/src/routes/messages/host.rs | 186 +++++++++++++++++ .../python-bridge/src/routes/messages/mod.rs | 64 ++++++ .../crates/python-bridge/src/routes/mod.rs | 31 --- .../python-bridge/src/routes/ocr/host.rs | 4 + litellm/rust_bridge/_native.pyi | 28 +-- litellm/rust_bridge/catalog.py | 1 + litellm/rust_bridge/failures.py | 43 ++++ litellm/rust_bridge/legacy_callbacks.py | 80 +++++++ litellm/rust_bridge/lifecycle.py | 137 ++++++++++-- litellm/rust_bridge/messages/entrypoints.py | 10 +- litellm/rust_bridge/messages/route_host.py | 2 +- litellm/rust_bridge/ocr/route_host.py | 40 +--- .../rust_bridge/native_route_wheel_test.py | 31 +-- .../test_litellm/rust_bridge/test_bindings.py | 4 +- .../test_litellm/rust_bridge/test_catalog.py | 4 + .../test_litellm/rust_bridge/test_runtime.py | 1 - tests/test_litellm_rust/messages/__init__.py | 0 .../messages/test_callbacks.py | 175 ++++++++++++++++ .../support/recording_server.py | 16 +- tests/test_litellm_rust/support/requests.py | 35 ++++ 40 files changed, 1319 insertions(+), 344 deletions(-) create mode 100644 litellm-rust/crates/core/src/messages/route.rs delete mode 100644 litellm-rust/crates/python-bridge/src/routes/messages.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/host.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/mod.rs create mode 100644 tests/test_litellm_rust/messages/__init__.py create mode 100644 tests/test_litellm_rust/messages/test_callbacks.py diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy/python_contract.json index 840c0abfa45..a09bdc711a3 100644 --- a/litellm-rust/crates/callbacks-legacy/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy/python_contract.json @@ -94,5 +94,22 @@ "kwargs", "error", "call_type" + ], + "stream_opened": [ + "logger" + ], + "stream_success": [ + "logger", + "request_body", + "chunks", + "start", + "end", + "first_chunk" + ], + "stream_failure": [ + "logger", + "request_body", + "chunks", + "error" ] } diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 53aa6ce9d2e..9d28db92add 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -2,7 +2,9 @@ //! raises is answered with the same `Logging` calls, in the same order, as the Python //! `@client` path makes them. -use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest}; +use litellm_callbacks::event::{ + CallEvent, FailureOrigin, RequestContext, Timing, WireRequest, epoch_seconds, +}; use litellm_host_python::{ LifecycleStep, PublicValue, PythonLifecycle, from_py, missing_state, to_py, }; @@ -10,14 +12,16 @@ use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::PyDict, + types::{PyDict, PyList}, }; use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, - finalize, is_internal_call, prepare, setup, + finalize, is_internal_call, + legacy_python::Streaming, + prepare, setup, }; /// What the legacy contract needs to know about the route it is logging. @@ -28,6 +32,12 @@ pub struct LegacySurface { pub input_description: &'static str, } +/// What the Messages stream iterator keeps for its end-of-stream billing. +struct DeliveredStream { + chunks: Py, + first_chunk: Option>, +} + enum Pending { DeploymentPreCall, DeploymentPostCall, @@ -46,6 +56,7 @@ pub struct LegacyLogging { body: Option>, headers: Option>, context: Option, + stream: Option, asynchronous: bool, internal: bool, pending: Option, @@ -80,6 +91,7 @@ impl LegacyLogging { body: None, headers: None, context: None, + stream: None, asynchronous, internal: false, pending: None, @@ -163,6 +175,49 @@ impl LegacyLogging { logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) } + fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> { + let logger = self.logger()?; + let billed = Streaming::Success.call( + py, + ( + logger.object(py), + &self.body, + &stream.chunks, + &self.start, + &self.end, + &stream.first_chunk, + ), + ); + match billed { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, Some(logger.object(py))); + Ok(()) + } + result => result.map(|_| ()), + } + } + + /// A failure after the stream reached the caller bills the delivered chunks as + /// partial usage. The sync path has no loop to schedule that on, so it falls back to + /// the plain failure handler. + fn stream_failure(&mut self, py: Python<'_>) -> PyResult { + let (Some(logger), Some(error), Some(stream)) = (&self.logger, &self.error, &self.stream) + else { + return Ok(LifecycleStep::Done); + }; + if !self.asynchronous { + return self.dispatch_failure(py); + } + match Streaming::Failure.call(py, (logger.object(py), &self.body, &stream.chunks, error)) { + Ok(awaitable) => { + self.pending = Some(Pending::AsyncFailure); + Ok(LifecycleStep::Await(awaitable.unbind())) + } + Err(failure) if is_cancellation(py, &failure) => Err(failure), + Err(_) => Ok(LifecycleStep::Done), + } + } + /// The sync failure handler, then the async one for async calls. Ordinary handler /// errors never replace the selected failure or suppress the other family; a /// cancellation does end the call. @@ -296,6 +351,22 @@ impl PythonLifecycle for LegacyLogging { ) -> PyResult { match (event, public) { (CallEvent::Started { .. }, _) => Ok(LifecycleStep::Done), + (CallEvent::Opened, _) => { + Streaming::Opened.call(py, (self.logger()?.object(py),))?; + self.stream = Some(DeliveredStream { + chunks: PyList::empty(py).unbind(), + first_chunk: None, + }); + Ok(LifecycleStep::Done) + } + (CallEvent::Delivered, Some(PublicValue::Chunk(chunk))) => { + let stream = self.stream.as_mut().ok_or_else(missing_state)?; + if stream.first_chunk.is_none() { + stream.first_chunk = Some(datetime(py, epoch_seconds())?); + } + stream.chunks.bind(py).append(chunk)?; + Ok(LifecycleStep::Done) + } (CallEvent::ResponseReceived { raw }, _) => { let api_key = self .context @@ -314,12 +385,18 @@ impl PythonLifecycle for LegacyLogging { (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); - self.dispatch_success(py)?; + match &self.stream { + Some(stream) => self.stream_success(py, stream)?, + None => self.dispatch_success(py)?, + } Ok(LifecycleStep::Done) } (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); + if self.stream.is_some() { + return self.stream_failure(py); + } if *origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() @@ -366,6 +443,7 @@ impl PythonLifecycle for LegacyLogging { } self.body = None; self.context = None; + self.stream = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -377,6 +455,10 @@ impl PythonLifecycle for LegacyLogging { visit.call(&self.end)?; visit.call(&self.response)?; visit.call(&self.error)?; + if let Some(stream) = &self.stream { + visit.call(&stream.chunks)?; + visit.call(&stream.first_chunk)?; + } visit.call(&self.body) } } diff --git a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs index a924775070c..7f5c77c1735 100644 --- a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs +++ b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs @@ -16,6 +16,7 @@ pub(crate) enum LegacyPython { Wrapper(Wrapper), Logging(Logging), DeploymentHooks(DeploymentHooks), + Streaming(Streaming), } /// The `@client` wrapper around the call: `function_setup`, limits, credentials, @@ -76,12 +77,25 @@ pub(crate) enum DeploymentHooks { AfterDeploymentFailure, } +/// The Messages stream iterator's logging: the stream flag, the end-of-stream billing +/// from the delivered chunks, and the partial-usage failure path. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Streaming { + #[strum(serialize = "stream_opened")] + Opened, + #[strum(serialize = "stream_success")] + Success, + #[strum(serialize = "stream_failure")] + Failure, +} + impl LegacyPython { fn name(self) -> &'static str { match self { Self::Wrapper(function) => function.into(), Self::Logging(function) => function.into(), Self::DeploymentHooks(function) => function.into(), + Self::Streaming(function) => function.into(), } } @@ -111,6 +125,15 @@ impl Logging { } } +impl Streaming { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Streaming(self).call(py, args) + } +} + impl DeploymentHooks { pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> where @@ -126,7 +149,7 @@ mod tests { use strum::VariantArray; - use super::{DeploymentHooks, LegacyPython, Logging, Wrapper}; + use super::{DeploymentHooks, LegacyPython, Logging, Streaming, Wrapper}; use crate::test_support::PYTHON_CONTRACT; #[test] @@ -147,6 +170,11 @@ mod tests { .iter() .map(|&function| LegacyPython::DeploymentHooks(function)), ) + .chain( + Streaming::VARIANTS + .iter() + .map(|&function| LegacyPython::Streaming(function)), + ) .map(LegacyPython::name) .collect(); assert_eq!(called.len(), declared.len(), "a function is borrowed twice"); diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy/tests/support.rs index 444655ea77b..42ca184eb16 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/support.rs @@ -84,6 +84,11 @@ FAKES = { 'success', response, call_type ), 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), + 'stream_opened': lambda logger: logger.record('stream_opened', None), + 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success', list(chunks) + ), + 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), } assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) for name, fake in FAKES.items(): diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs index 3bf12e553b9..d19e973b812 100644 --- a/litellm-rust/crates/callbacks/src/event.rs +++ b/litellm-rust/crates/callbacks/src/event.rs @@ -60,6 +60,10 @@ pub enum CallEvent { ResponseReceived { raw: RawResponse, }, + /// The call streams and its stream was handed to the caller. + Opened, + /// One chunk of an open stream reached the caller. + Delivered, Succeeded { timing: Timing, }, diff --git a/litellm-rust/crates/callbacks/src/host.rs b/litellm-rust/crates/callbacks/src/host.rs index 2392718a18d..eef3e1da8d5 100644 --- a/litellm-rust/crates/callbacks/src/host.rs +++ b/litellm-rust/crates/callbacks/src/host.rs @@ -11,12 +11,25 @@ pub enum HostOp { context: Box, }, Emit(CallEvent), + /// The response streams: the host hands the caller a stream and answers once the + /// caller asks for the first chunk or goes away. + Open(R::StreamHead), + /// The next chunk of an open stream, answered once the caller asks for the one after. + Deliver(R::Chunk), } pub enum HostResult { Route(R::OpResult), BeforeSend(Box), Emitted, + Demand(Demand), +} + +/// Whether the caller of a streamed call still reads it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Demand { + More, + Detached, } /// A host answer that is either available now or arrives once the host's own @@ -42,4 +55,12 @@ pub trait Host: Send + Sync { fn emit(&self, _event: &CallEvent) -> impl Future> + Send { async { Ok(()) } } + + fn open(&self, _head: R::StreamHead) -> impl Future> + Send { + async { Ok(Demand::More) } + } + + fn deliver(&self, _chunk: R::Chunk) -> impl Future> + Send { + async { Ok(Demand::More) } + } } diff --git a/litellm-rust/crates/callbacks/src/route.rs b/litellm-rust/crates/callbacks/src/route.rs index 97738c8da8b..8ab2b125760 100644 --- a/litellm-rust/crates/callbacks/src/route.rs +++ b/litellm-rust/crates/callbacks/src/route.rs @@ -6,4 +6,9 @@ pub trait Route: Send + Sync + 'static { type Error: Clone + Send + Sync + 'static; type Op: Send + 'static; type OpResult: Send + 'static; + /// One piece of a streamed response, handed to the caller as it arrives. A route + /// that never streams uses `Infallible`. + type Chunk: Send + 'static; + /// What the route knows once a streamed response starts, before its first chunk. + type StreamHead: Send + 'static; } diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/callbacks/src/run.rs index 5accaa2e25e..705c5504c01 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/callbacks/src/run.rs @@ -26,6 +26,8 @@ where .await .map(|wire| HostResult::BeforeSend(Box::new(wire))), HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted), + HostOp::Open(head) => host.open(head).await.map(HostResult::Demand), + HostOp::Deliver(chunk) => host.deliver(chunk).await.map(HostResult::Demand), }; match answer { Ok(answer) => result = Some(answer), @@ -61,6 +63,8 @@ mod tests { type Error = &'static str; type Op = &'static str; type OpResult = (); + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } struct Scripted { diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/core/src/machine/mod.rs index 279a2d65c97..929f0a423c4 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/core/src/machine/mod.rs @@ -9,7 +9,7 @@ use std::{future::Future, pin::Pin}; pub use auth::{HostTokenProvider, TokenRoute}; use litellm_callbacks::{ event::{CallEvent, RequestContext, WireRequest}, - host::{HostOp, HostResult}, + host::{Demand, HostOp, HostResult}, machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, route::Route, }; @@ -88,6 +88,21 @@ where _ => Err(MachineFault::Mismatch.into()), } } + + pub async fn open(&self, head: R::StreamHead) -> Result { + self.demand(HostOp::Open(head)).await + } + + pub async fn deliver(&self, chunk: R::Chunk) -> Result { + self.demand(HostOp::Deliver(chunk)).await + } + + async fn demand(&self, op: HostOp) -> Result { + match self.invoke(op).await? { + HostResult::Demand(demand) => Ok(demand), + _ => Err(MachineFault::Mismatch.into()), + } + } } enum Execution { diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index b95402b1a7a..22e2c398ff7 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,88 +1,54 @@ -use litellm_llms::custom_httpx::http_handler::http_request; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use std::time::Duration; -use super::{ - Error, client::http_client, common_utils::truncate_error_body, - prepare::prepare_provider_request, +use litellm_llms::{ + base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, + custom_httpx::{http_handler::http_request, transport::Error as TransportError}, }; -use crate::{constants::ANTHROPIC_MESSAGES_PROVIDER, messages::types::MessagesRequest}; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use serde_json::Value; -pub(super) async fn execute_messages_provider_call( - request: MessagesRequest<'_>, +use super::{Error, client::http_client, common_utils::truncate_error_body}; + +pub(super) fn network(error: reqwest::Error) -> Error { + Error::Transport(TransportError::Network(error.to_string())) +} + +pub(super) async fn send( + url: &str, + headers: &[(String, String)], + body: &Value, + timeout: Option, +) -> Result { + let builder = headers.iter().fold( + http_client().post(url).json(body), + |builder, (key, value)| builder.header(key, value), + ); + let builder = match timeout { + Some(duration) => builder.timeout(duration), + None => builder, + }; + http_request(builder).await.map_err(network) +} + +pub(super) async fn provider_error(response: reqwest::Response) -> Error { + let status = response.status().as_u16(); + match response.text().await { + Ok(text) => Error::Transport(TransportError::Http { + status, + body: truncate_error_body(&text), + }), + Err(error) => network(error), + } +} + +pub(super) fn decode_response( + config: &dyn BaseAnthropicMessagesConfig, + model: &str, + text: &str, ) -> Result { - let request = prepare_provider_request(request)?; - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder).await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - - let status = response.status(); - let text = response.text().await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - - if !status.is_success() { - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }, - )); - } - - let response = serde_json::from_str(&text) + let response = serde_json::from_str(text) .map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?; - request - .config - .transform_anthropic_messages_response(&request.model, response) + config + .transform_anthropic_messages_response(model, response) .map_err(Error::from) } - -pub(super) async fn execute_messages_provider_stream( - request: MessagesRequest<'_>, -) -> Result { - let request = prepare_provider_request(request)?; - if request.provider != ANTHROPIC_MESSAGES_PROVIDER { - return Err(Error::Unsupported("streaming messages for this provider")); - } - - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder).await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - let status = response.status(); - if !status.is_success() { - let text = response.text().await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }, - )); - } - Ok(response) -} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index c3d7bea48ff..e36c6668efe 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -1,11 +1,8 @@ //! The Anthropic Messages call, the Rust equivalent of Python's //! `litellm.messages()`. //! -//! [`messages`] is the top-level entrypoint: give it a model, a body, and -//! credentials, and it resolves the provider, transforms the request, calls the -//! provider, and returns a typed non-streaming response. [`messages_stream`] -//! is the streaming variant; it hands the raw upstream response back so a host -//! can splice the event stream to its own caller. +//! [`route`] is the call as a machine a host drives, streaming or not. [`messages`] runs +//! it in process for a caller that already holds the request and wants the message. mod error; pub mod types; @@ -14,17 +11,34 @@ mod client; mod common_utils; mod handler; mod prepare; -use handler::{execute_messages_provider_call, execute_messages_provider_stream}; +pub mod route; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}; +use serde_json::Value; use crate::messages::types::MessagesRequest; pub async fn messages(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_call(request).await -} - -pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_stream(request).await + let Value::Object(body) = request.body else { + return Err(Error::InvalidRequest( + "messages body must be an object".into(), + )); + }; + let call = MessagesCall { + model: request.model.into(), + body, + api_key: request.api_key.map(Into::into), + api_base: request.api_base.map(Into::into), + custom_llm_provider: request.custom_llm_provider.map(Into::into), + extra_headers: request.extra_headers, + timeout: request.timeout, + }; + match litellm_callbacks::run::run(messages_machine(), &LocalMessagesHost::new(call)).await? { + MessagesOutput::Message(message) => Ok(message), + MessagesOutput::Streamed => Err(Error::Unsupported( + "streamed responses need a streaming host", + )), + } } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 8b676803871..850f9108869 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -2,6 +2,7 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_l use litellm_llms::base_llm::anthropic_messages::transformation::{ BaseAnthropicMessagesConfig, MessagesAuthStrategy, }; +use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use serde_json::{Map, Value}; use super::{ @@ -37,10 +38,14 @@ pub(super) fn prepare_provider_request( let headers = validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?; - let typed_request = serde_json::from_value(request.body).map_err(|err| { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + let typed_request: AnthropicMessagesRequest = + serde_json::from_value(request.body).map_err(|err| { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + })?; + let transformed = config.transform_anthropic_messages_request(AnthropicMessagesRequest { + model: model.clone(), + ..typed_request })?; - let transformed = config.transform_anthropic_messages_request(typed_request)?; let body = serde_json::to_value(transformed).map_err(|err| { Error::InvalidRequest(format!( "failed to serialize Anthropic messages request: {err}" diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs new file mode 100644 index 00000000000..a680b5ce1ec --- /dev/null +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -0,0 +1,197 @@ +use std::{sync::Mutex, time::Duration}; + +use bytes::Bytes; +use litellm_auth::SecretValue; +use litellm_callbacks::{ + event::{CallEvent, RawResponse, RequestContext, WireRequest}, + host::{Demand, Host}, + route::Route, +}; +use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use serde_json::{Map, Value}; + +use super::{ + Error, + common_utils::messages_provider_config, + handler::{decode_response, network, provider_error, send}, + prepare::prepare_provider_request, + types::MessagesRequest, +}; +use crate::{ + constants::ANTHROPIC_MESSAGES_PROVIDER, + machine::{HostChannel, MachineFault, RouteMachine}, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MessagesOp { + ProjectRequest, +} + +pub enum MessagesOpResult { + Request(Box), +} + +/// The caller's request as the host projects it. +pub struct MessagesCall { + pub model: String, + pub body: Map, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, + pub timeout: Option, +} + +impl MessagesCall { + fn streams(&self) -> bool { + self.body.get("stream").and_then(Value::as_bool) == Some(true) + } +} + +pub enum MessagesOutput { + Message(AnthropicMessagesResponse), + /// Every chunk already reached the host through `Deliver`. + Streamed, +} + +pub struct Messages; + +impl Route for Messages { + type Response = MessagesOutput; + type Error = Error; + type Op = MessagesOp; + type OpResult = MessagesOpResult; + type Chunk = Bytes; + type StreamHead = (); +} + +impl From for Error { + fn from(fault: MachineFault) -> Self { + Self::InvalidRequest(match fault { + MachineFault::Abandoned => "messages host driver was abandoned".into(), + MachineFault::Protocol(message) => format!("messages {message}"), + MachineFault::Mismatch => "invalid messages host operation result".into(), + }) + } +} + +pub type MessagesHost = HostChannel; +pub type MessagesMachine = RouteMachine; + +/// Whether this route serves the request, decided before any callback runs so a host +/// can still run its own path. +pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool { + let provider = get_custom_llm_provider(model, custom_llm_provider) + .map(|resolved| resolved.custom_llm_provider) + .or(custom_llm_provider); + match provider { + Some(ANTHROPIC_MESSAGES_PROVIDER) => true, + Some(provider) => !stream && messages_provider_config(provider).is_some(), + None => false, + } +} + +/// The in-process host for a request already in hand. It answers projection once and +/// observes nothing. +pub struct LocalMessagesHost { + call: Mutex>, +} + +impl LocalMessagesHost { + pub fn new(call: MessagesCall) -> Self { + Self { + call: Mutex::new(Some(call)), + } + } +} + +impl Host for LocalMessagesHost { + async fn route(&self, op: MessagesOp) -> Result { + match op { + MessagesOp::ProjectRequest => self + .call + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .map(|call| MessagesOpResult::Request(Box::new(call))) + .ok_or_else(|| { + Error::InvalidRequest("messages request was already projected".into()) + }), + } + } +} + +pub fn messages_machine() -> MessagesMachine { + RouteMachine::new(|host| Box::pin(execute(host))) +} + +async fn execute(host: MessagesHost) -> Result { + let MessagesOpResult::Request(call) = host.route(MessagesOp::ProjectRequest).await?; + let stream = call.streams(); + let request = prepare_provider_request(MessagesRequest { + model: &call.model, + body: Value::Object(call.body.clone()), + api_key: call.api_key.as_deref(), + api_base: call.api_base.as_deref(), + custom_llm_provider: call.custom_llm_provider.as_deref(), + extra_headers: call.extra_headers.clone(), + timeout: call.timeout, + })?; + if stream && request.provider != ANTHROPIC_MESSAGES_PROVIDER { + return Err(Error::Unsupported("streaming messages for this provider")); + } + let context = RequestContext { + model: request.model.clone(), + custom_llm_provider: request.provider.clone(), + optional_params: Value::Object( + call.body + .iter() + .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + ), + secret_fields: Vec::new(), + api_key: call.api_key.clone().map(SecretValue::new), + }; + let wire = host + .before_send( + WireRequest { + url: request.url, + headers: request.upstream_headers, + body: request.body, + }, + context, + ) + .await?; + let response = send(&wire.url, &wire.headers, &wire.body, request.timeout).await?; + if !response.status().is_success() { + return Err(provider_error(response).await); + } + if stream { + return relay(&host, response).await; + } + let text = response.text().await.map_err(network)?; + host.emit(CallEvent::ResponseReceived { + raw: RawResponse { body: text.clone() }, + }) + .await?; + decode_response(request.config, &request.model, &text).map(MessagesOutput::Message) +} + +/// Hands each upstream chunk to the caller as it arrives. A caller that stops reading +/// ends the upstream read, and the call completes with what it delivered. +async fn relay( + host: &MessagesHost, + mut response: reqwest::Response, +) -> Result { + if host.open(()).await? == Demand::Detached { + return Ok(MessagesOutput::Streamed); + } + while let Some(chunk) = response.chunk().await.map_err(network)? { + if host.deliver(chunk).await? == Demand::Detached { + break; + } + } + Ok(MessagesOutput::Streamed) +} diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index ac4237651da..e6ee45c64a8 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -39,6 +39,8 @@ impl Route for Ocr { type Error = Error; type Op = OcrOp; type OpResult = OcrOpResult; + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } impl TokenRoute for Ocr { diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index ca2e14a7f0d..779d2037bb3 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -198,6 +198,8 @@ fn event_name(event: &CallEvent) -> &'static str { CallEvent::ResponseReceived { .. } => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", + CallEvent::Opened => "opened", + CallEvent::Delivered => "delivered", } } diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 4aa7a2163ca..795977438fd 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -23,6 +23,7 @@ pub enum LifecycleStep { pub enum PublicValue<'a> { Response(&'a Py), Error(&'a PyErr), + Chunk(&'a Py), } /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in @@ -111,6 +112,13 @@ pub trait RouteHost: Send + Sync { response: ::Response, ) -> PyResult>; + /// One streamed chunk as the caller receives it. + fn chunk( + &mut self, + py: Python<'_>, + chunk: ::Chunk, + ) -> PyResult>; + fn classify( &self, py: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 7895d803cb5..c59ba1925dc 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -3,7 +3,7 @@ use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing, epoch_seconds}; -use litellm_callbacks::host::{HostOp, HostResult, HostStep}; +use litellm_callbacks::host::{Demand, HostOp, HostResult, HostStep}; use litellm_callbacks::machine::{HostFailure, Machine, MachineStep}; use litellm_callbacks::route::Route; use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; @@ -38,6 +38,7 @@ struct MachineState { enum Stage { Begin, Call, + Streaming, AfterSuccess, Succeeded(Py), Failed(Py), @@ -56,6 +57,8 @@ enum Expect { enum Pending { Native, Adapter(Expect), + /// The stream handed to the caller waits for its next read or its close. + Consumer, } enum Next { @@ -121,7 +124,14 @@ where } match driver.resume(None)? { ExecutionStep::Return(value) => Ok(value), - ExecutionStep::Await(_) => Err(PyRuntimeError::new_err("sync call suspended")), + ExecutionStep::Open => py + .import("litellm.rust_bridge.lifecycle")? + .getattr("SyncStream")? + .call1((Py::new(py, Execution::suspended(driver))?,)) + .map(Bound::unbind), + ExecutionStep::Await(_) | ExecutionStep::Yield(_) => { + Err(PyRuntimeError::new_err("sync call suspended")) + } } } @@ -162,6 +172,14 @@ where self.run_steps(py, HostStep::Ready(result)) } (Some(Pending::Native), Some(Err(error))) => self.interrupt(py, error), + (Some(Pending::Consumer), Some(read)) => { + let demand = if read.is_ok() { + Demand::More + } else { + Demand::Detached + }; + self.resume_machine(py, Some(Ok(HostResult::Demand(demand)))) + } (Some(Pending::Adapter(expect)), Some(result)) => { match self.adapter.resume(py, result) { Ok(step) => self.on_adapter(py, step, expect), @@ -216,7 +234,7 @@ where fn adapter_failed(&mut self, py: Python<'_>, error: PyErr) -> PyResult { match self.stage { Stage::Begin | Stage::AfterSuccess => self.failure(py, error, FailureOrigin::Host), - Stage::Call => self.interrupt(py, error), + Stage::Call | Stage::Streaming => self.interrupt(py, error), Stage::Succeeded(_) | Stage::Failed(_) => Err(error), } } @@ -283,6 +301,8 @@ where Err(error) => Err(error), } } + HostOp::Open(_) => return self.opened(py).map(Next::Return), + HostOp::Deliver(chunk) => return self.delivered(py, chunk).map(Next::Return), HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), Ok(LifecycleStep::Await(awaitable)) => { @@ -299,6 +319,40 @@ where } } + fn opened(&mut self, py: Python<'_>) -> PyResult { + self.stage = Stage::Streaming; + match self.adapter.emit(py, &CallEvent::Opened, None) { + Ok(LifecycleStep::Done) => { + self.pending = Some(Pending::Consumer); + Ok(ExecutionStep::Open) + } + Ok(_) => Err(missing_state()), + Err(error) => self.interrupt(py, error), + } + } + + fn delivered( + &mut self, + py: Python<'_>, + chunk: as Route>::Chunk, + ) -> PyResult { + let chunk = match self.route.chunk(py, chunk) { + Ok(chunk) => chunk, + Err(error) => return self.interrupt(py, error), + }; + let observed = + self.adapter + .emit(py, &CallEvent::Delivered, Some(PublicValue::Chunk(&chunk))); + match observed { + Ok(LifecycleStep::Done) => { + self.pending = Some(Pending::Consumer); + Ok(ExecutionStep::Yield(chunk)) + } + Ok(_) => Err(missing_state()), + Err(error) => self.interrupt(py, error), + } + } + fn interrupt(&mut self, py: Python<'_>, error: PyErr) -> PyResult { let cancelled = is_cancellation(py, &error); let native = H::host_error(&error); @@ -369,6 +423,9 @@ where Ok(public) => public, Err(error) => return self.failure(py, error, FailureOrigin::Call), }; + if let Stage::Streaming = self.stage { + return self.succeeded(py, public); + } self.stage = Stage::AfterSuccess; match self.adapter.after_success(py, public, self.timing()) { Ok(step) => self.on_adapter(py, step, Expect::Response), @@ -536,6 +593,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri type Error = Error; type Op = &'static str; type OpResult = String; + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } /// Yields the scripted ops in order, then completes or fails as scripted. @@ -574,6 +633,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri HostResult::Route(value) => value, HostResult::BeforeSend(wire) => wire.url, HostResult::Emitted => "emitted".into(), + HostResult::Demand(demand) => format!("{demand:?}"), }); } if !self.ops.is_empty() { @@ -648,6 +708,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { + match chunk {} + } + fn complete(&mut self, py: Python<'_>, response: String) -> PyResult> { self.log.push("complete"); Ok(pyo3::types::PyString::new(py, &response) @@ -1138,6 +1202,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ) .into()) } + fn chunk( + &mut self, + _: Python<'_>, + chunk: std::convert::Infallible, + ) -> PyResult> { + match chunk {} + } fn complete(&mut self, _: Python<'_>, _: String) -> PyResult> { Err(missing_state()) } diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index d8cd6c92130..10abbadbda5 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -8,6 +8,10 @@ use pyo3::prelude::*; pub enum ExecutionStep { Return(Py), Await(Py), + /// The call streams: the caller gets a stream over this execution, which stays + /// suspended until the stream asks for a chunk. + Open, + Yield(Py), } pub trait ExecutionBody: Send + Sync { @@ -34,6 +38,13 @@ impl Execution { } } + /// An execution already started elsewhere and now waiting for its next input. + pub fn suspended(body: impl ExecutionBody + 'static) -> Self { + Self { + state: ExecutionState::Suspended(Box::new(body)), + } + } + fn advance( slf: &Bound<'_, Self>, py: Python<'_>, @@ -64,6 +75,8 @@ impl Execution { let step = body.resume(result)?; let (tag, value, suspended) = match step { ExecutionStep::Await(value) => ("Await", value, true), + ExecutionStep::Open => ("Open", py.None(), true), + ExecutionStep::Yield(value) => ("Yield", value, true), ExecutionStep::Return(value) => ("Complete", value, false), }; let step = py diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index ea4077b102f..2aba51cc4ff 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -18,10 +18,6 @@ pub(crate) struct RouteOptions { pub(crate) timeout: Option, } -pub(crate) fn body_argument(value: &Bound<'_, PyAny>) -> PyResult> { - required_object("body", from_py_argument(value)?) -} - pub(crate) fn messages_argument(value: &Bound<'_, PyAny>) -> PyResult> { match from_py_argument(value)? { Value::Array(values) => Ok(values), @@ -192,18 +188,6 @@ mod tests { json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) ); - let body = py - .eval( - c"{'model': 'claude', 'metadata': {'user': '1'}}", - None, - None, - ) - .unwrap(); - assert_eq!( - Value::Object(body_argument(&body).unwrap()), - json!({"model": "claude", "metadata": {"user": "1"}}) - ); - let params = py.eval(c"{'temperature': 0.2}", None, None).unwrap(); assert_eq!( optional_params_argument(¶ms).unwrap(), diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs deleted file mode 100644 index daec931c92e..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ /dev/null @@ -1,88 +0,0 @@ -use litellm_core::messages::{Error, messages as run_messages, types::MessagesRequest}; -use litellm_host_python::{run_async, run_sync}; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; -use pyo3::prelude::*; -use serde_json::{Map, Value}; - -use crate::{ - errors::messages_error_to_pyerr, - marshal::{RouteOptions, body_argument, extra_headers_argument, optional_timeout}, -}; - -async fn execute( - body: Map, - options: RouteOptions, -) -> Result { - let RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout, - } = options; - run_messages(MessagesRequest { - model: &model, - body: Value::Object(body), - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[expect( - clippy::too_many_arguments, - reason = "one parameter per Python keyword" -)] -pub(crate) fn messages( - py: Python<'_>, - model: String, - #[pyo3(from_py_with = body_argument)] body: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - #[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let options = RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout: optional_timeout(timeout_seconds), - }; - run_sync(py, execute(body, options), messages_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[expect( - clippy::too_many_arguments, - reason = "one parameter per Python keyword" -)] -pub(crate) fn amessages<'py>( - py: Python<'py>, - model: String, - #[pyo3(from_py_with = body_argument)] body: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - #[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let options = RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout: optional_timeout(timeout_seconds), - }; - run_async(py, execute(body, options), messages_error_to_pyerr) -} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs new file mode 100644 index 00000000000..0c87aea1ca2 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -0,0 +1,186 @@ +use bytes::Bytes; +use litellm_core::messages::{ + Error, + route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput}, +}; +use litellm_host_python::{HostOpError, RouteHost, from_py, lookup, to_py}; +use litellm_llms::custom_httpx::transport::Error as TransportError; +use pyo3::{ + exceptions::{PyException, PyValueError}, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::{PyBytes, PyDict}, +}; +use serde_json::{Map, Value}; + +use crate::{ + errors::{RustUpstreamError, messages_error_to_pyerr}, + marshal::{optional_timeout, python_timeout_seconds}, +}; + +/// The Anthropic Messages body fields a caller may pass besides `model` and `messages`, +/// as `AnthropicMessagesRequestOptionalParams` declares them. +const BODY_FIELDS: [&str; 20] = [ + "max_tokens", + "metadata", + "stop_sequences", + "stream", + "system", + "temperature", + "thinking", + "tool_choice", + "tools", + "top_k", + "inference_geo", + "top_p", + "mcp_servers", + "context_management", + "container", + "output_format", + "speed", + "output_config", + "cache_control", + "reasoning_effort", +]; + +/// The Python side of the Messages route: projects the prepared arguments and builds the +/// public response, chunks and exceptions. +pub(super) struct MessagesRouteHost { + request: Py, +} + +impl MessagesRouteHost { + pub(super) fn new(request: Py) -> Self { + Self { request } + } + + fn project(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { + let request = self.request.bind(py); + let argument = |name: &str| -> PyResult>> { + Ok(lookup(arguments, request, name)?.filter(|value| !value.is_none())) + }; + let string = |name: &str| -> PyResult> { + argument(name)?.map(|value| value.extract()).transpose() + }; + let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?; + let messages = + argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?; + let fields = BODY_FIELDS + .iter() + .filter_map(|name| match argument(name) { + Ok(Some(value)) => Some(from_py(&value).map(|value| ((*name).to_string(), value))), + Ok(None) => None, + Err(error) => Some(Err(error)), + }) + .collect::>>()?; + let body = [ + ("model".to_string(), Value::String(model.clone())), + ("messages".to_string(), from_py(&messages)?), + ] + .into_iter() + .chain(fields) + .collect::>(); + let timeout = argument("timeout")? + .map(|value| python_timeout_seconds(py, value.unbind())) + .transpose()? + .flatten(); + Ok(MessagesCall { + model, + body, + api_key: string("api_key")?, + api_base: string("api_base")?, + custom_llm_provider: string("custom_llm_provider")?, + extra_headers: argument("extra_headers")? + .map(|value| from_py(&value)) + .transpose()?, + timeout: optional_timeout(timeout), + }) + } + + fn provider(&self, py: Python<'_>) -> String { + self.request + .bind(py) + .getattr("custom_llm_provider") + .and_then(|value| value.extract::>()) + .ok() + .flatten() + .unwrap_or_else(|| "anthropic".into()) + } + + fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { + if !error.is_instance_of::(py) { + return error; + } + let mapped = py + .import("litellm.rust_bridge.messages.route_host") + .and_then(|module| module.getattr("map_failure")) + .and_then(|map| map.call1((error.value(py), self.request.bind(py), self.provider(py)))) + .and_then(|mapped| { + mapped + .extract::>() + .map_err(PyErr::from) + }); + match mapped { + Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), + Err(_) => error, + } + } +} + +impl RouteHost for MessagesRouteHost { + type Route = Messages; + type Failure = PyErr; + + fn invoke( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + op: MessagesOp, + ) -> Result> { + match op { + MessagesOp::ProjectRequest => self + .project(py, arguments) + .map(|call| MessagesOpResult::Request(Box::new(call))) + .map_err(|error| HostOpError::Python(self.map_failure(py, error))), + } + } + + fn complete(&mut self, py: Python<'_>, response: MessagesOutput) -> PyResult> { + match response { + MessagesOutput::Message(message) => py + .import("litellm.rust_bridge.messages.route_host")? + .getattr("response")? + .call1((to_py(py, &message)?,)) + .map(Bound::unbind), + MessagesOutput::Streamed => Ok(py.None()), + } + } + + fn chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { + Ok(PyBytes::new(py, &chunk).into_any().unbind()) + } + + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + let native = match error { + Error::Transport(TransportError::Http { status, body }) => { + let error = RustUpstreamError::new_err((status, body)); + error + .value(py) + .setattr("headers", Vec::<(String, String)>::new())?; + error + } + other => messages_error_to_pyerr(other), + }; + Ok(self.map_failure(py, native)) + } + + fn host_error(error: &PyErr) -> Error { + Error::InvalidRequest(error.to_string()) + } + + fn close(&mut self, _: Python<'_>) {} + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.request) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs new file mode 100644 index 00000000000..b4259aca5ba --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -0,0 +1,64 @@ +mod host; + +use host::MessagesRouteHost; +use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; +use litellm_core::messages::route::{messages_machine, supports}; +use pyo3::{ + prelude::*, + types::{PyDict, PyTuple}, +}; + +use crate::errors::RustBridgeDeclined; + +const SURFACE: LegacySurface = LegacySurface { + call_type: "anthropic_messages", + input_description: "Messages", +}; + +fn run_messages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, +) -> PyResult> { + let model: String = request.getattr("model")?.extract()?; + let provider: Option = request.getattr("custom_llm_provider")?.extract()?; + let stream = request + .getattr("stream")? + .extract::>()? + .unwrap_or(false); + if !supports(&model, provider.as_deref(), stream) { + return Err(RustBridgeDeclined::new_err( + "the Rust Messages route does not serve this provider", + )); + } + run_legacy_call( + py, + SURFACE, + PublicCall::capture(&request, &args, &kwargs)?, + messages_machine(), + MessagesRouteHost::new(request.unbind()), + asynchronous, + ) +} + +#[pyfunction] +pub(crate) fn messages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, +) -> PyResult> { + run_messages(py, request, args, kwargs, false) +} + +#[pyfunction] +pub(crate) fn amessages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, +) -> PyResult> { + run_messages(py, request, args, kwargs, true) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index f59e32a28e2..2d6b849a6b1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -22,11 +22,6 @@ mod tests { "atranscription", "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", ), - ( - "messages", - "amessages", - "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)", - ), ( "chat_completions", "achat_completions", @@ -113,25 +108,6 @@ value = Broken() ); assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string()); - let invalid_body = PyList::empty(py); - let sync_messages_error = module - .getattr("messages") - .and_then(|function| function.call1(("model", &invalid_body))) - .expect_err("sync Messages should reject a non-dict body"); - let async_messages_error = module - .getattr("amessages") - .and_then(|function| function.call1(("model", &invalid_body))) - .expect_err("async Messages should reject a non-dict body"); - - assert_eq!( - sync_messages_error.to_string(), - "ValueError: body must be a dict" - ); - assert_eq!( - async_messages_error.to_string(), - sync_messages_error.to_string() - ); - let invalid_headers = PyList::empty(py); let kwargs = PyDict::new(py); kwargs @@ -193,13 +169,6 @@ value = Broken() headers_kwargs .set_item("extra_headers", &invalid) .expect("kwargs should accept extra_headers"); - let invalid_body = PyList::empty(py); - let error = module - .getattr("messages") - .and_then(|function| function.call(("model", &invalid_body), Some(&headers_kwargs))) - .expect_err("body should be validated before headers"); - assert_eq!(error.to_string(), "ValueError: body must be a dict"); - let invalid_payload = PyModule::new(py, "invalid_payload").expect("invalid payload should be created"); let error = module diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 77212e3d38e..19211671fd7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -125,6 +125,10 @@ impl RouteHost for OcrRouteHost { .map(Bound::unbind) } + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { + match chunk {} + } + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { Ok(self.map_failure(py, ocr_error_to_pyerr(error))) } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 488e278cca7..9f959c056de 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,9 +1,11 @@ from asyncio import Future -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from typing import Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse class RustBridgeDeclined(Exception): ... class RustUpstreamError(Exception): ... @@ -39,23 +41,15 @@ def atranscription( timeout_seconds: float | None = None, ) -> Future[dict[str, object]]: ... def messages( - model: str, - body: Mapping[str, object], - api_key: str | None = None, - api_base: str | None = None, - custom_llm_provider: str | None = None, - extra_headers: Mapping[str, object] | None = None, - timeout_seconds: float | None = None, -) -> dict[str, object]: ... + request: LiteLLMMessagesRequest, + args: tuple[object, ...], + kwargs: dict[str, object], +) -> AnthropicMessagesResponse | Iterator[bytes]: ... def amessages( - model: str, - body: Mapping[str, object], - api_key: str | None = None, - api_base: str | None = None, - custom_llm_provider: str | None = None, - extra_headers: Mapping[str, object] | None = None, - timeout_seconds: float | None = None, -) -> Future[dict[str, object]]: ... + request: LiteLLMMessagesRequest, + args: tuple[object, ...], + kwargs: dict[str, object], +) -> Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[bytes]]: ... def chat_completions_decline( model: str, messages: Sequence[object], diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 9efbbfa2e9e..d843a874fe3 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -59,6 +59,7 @@ Rules: TypeAlias = tuple[Rule, ...] RULES: Final[Rules] = ( Rule(Route.OCR, Rollout.RUST_OPT_OUT), + Rule(Route.MESSAGES, Rollout.RUST_OPT_IN), Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/litellm/rust_bridge/failures.py b/litellm/rust_bridge/failures.py index b714341fe43..80805b7ff69 100644 --- a/litellm/rust_bridge/failures.py +++ b/litellm/rust_bridge/failures.py @@ -5,8 +5,37 @@ from __future__ import annotations from collections.abc import Mapping from typing import Final, Protocol, cast # noqa: TID251 # adapts the public exception mapper +import httpx +import openai +from pydantic import TypeAdapter, ValidationError + import litellm +_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) +_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) + + +class UpstreamFailure(Exception): + def __init__(self, response: httpx.Response, cause: Exception) -> None: + super().__init__(str(cause)) + self.message: Final = str(cause) + self.response: Final = response + self.status_code: Final = response.status_code + self.__cause__ = cause + + +def _upstream_failure(error: Exception, api_base: str | None) -> Exception: + try: + status, body = _UPSTREAM_ARGS.validate_python(error.args) + headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) + except ValidationError: + return error + http_request: Final = httpx.Request("POST", api_base or "https://docs.litellm.ai/docs") + return UpstreamFailure( + httpx.Response(status, content=body.encode(), headers=headers, request=http_request), + error, + ) + class ExceptionMapper(Protocol): def __call__( @@ -35,3 +64,17 @@ def map_failure(error: Exception, model: str, request_provider: str, kwargs: Map except Exception as public_error: public_error.__context__ = error return public_error + + +def map_native_failure( + error: Exception, model: str, request_provider: str, kwargs: Mapping[str, object], api_base: str | None = None +) -> Exception: + """`map_failure`, reading a native `(status, body)` provider failure as the HTTP response it was.""" + original: Final = _upstream_failure(error, api_base) + public_error: Final = map_failure(original, model, request_provider, kwargs) + if isinstance(original, UpstreamFailure) and public_error.__context__ is original: + public_error.__context__ = error + if isinstance(public_error, openai.APIStatusError): + public_error.response = original.response + public_error.status_code = original.status_code + return public_error diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index 406dc55cfea..fd0fac1799a 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -6,6 +6,7 @@ registries it fans out to. It expires with that contract. from __future__ import annotations +import asyncio import contextvars import datetime import traceback @@ -160,6 +161,21 @@ class LoggingWorker(Protocol): def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... +class StreamingLogBuilder(Protocol): + def __call__( + self, + *, + litellm_logging_obj: Logging, + passthrough_success_handler_obj: object, + url_route: str, + request_body: dict[str, object], + endpoint_type: object, + start_time: datetime.datetime, + raw_bytes: list[bytes], + end_time: datetime.datetime, + ) -> Coroutine[object, object, None]: ... + + class DeploymentHook(Protocol): def __call__(self, kwargs: dict[str, object], call_type: str) -> Awaitable[object]: ... @@ -306,3 +322,67 @@ def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_t DeploymentFailureHook, utils.async_post_call_failure_deployment_hook ) return hook(kwargs, error, call_type) + + +def stream_opened(logger: Logging) -> None: + logger.stream = True + logger.model_call_details["stream"] = True + + +def stream_success( + logger: Logging, + request_body: dict[str, object], + chunks: list[bytes], + start: datetime.datetime, + end: datetime.datetime, + first_chunk: datetime.datetime | None, +) -> None: + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + if first_chunk is not None: + logger.completion_start_time = first_chunk + logger.model_call_details["completion_start_time"] = first_chunk + build: Final = cast( # cast-ok: bounded adapter for the untyped pass-through logging builder + StreamingLogBuilder, + PassThroughStreamingHandler._route_streaming_logging_to_handler, # pyright: ignore[reportPrivateUsage] # the Messages stream iterator bills through the same builder + ) + coroutine: Final = build( + litellm_logging_obj=logger, + passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + url_route="/v1/messages", + request_body=request_body, + endpoint_type=EndpointType.ANTHROPIC, + start_time=start, + raw_bytes=chunks, + end_time=end, + ) + if getattr(logger, "_on_deferred_stream_complete", None) is not None: + logger._deferred_stream_complete_args = (coroutine,) # pyright: ignore[reportAttributeAccessIssue] # the proxy's deferred stream release reads this slot + return + try: + asyncio.get_running_loop() + except RuntimeError: + from litellm.litellm_core_utils.litellm_logging import executor + + executor.submit(contextvars.copy_context().run, asyncio.run, coroutine) + return + enqueue_logging(coroutine) + + +def stream_failure( + logger: Logging, request_body: dict[str, object], chunks: list[bytes], error: Exception +) -> Coroutine[object, object, None]: + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + return PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=logger, + endpoint_type=EndpointType.ANTHROPIC, + request_body=request_body, + raw_bytes=chunks, + exception=error, + ) diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index d903021b6f3..4096d386964 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -1,8 +1,8 @@ from __future__ import annotations -from collections.abc import Awaitable +from collections.abc import AsyncIterator, Awaitable, Iterator from dataclasses import dataclass -from typing import Protocol +from typing import Final, Protocol @dataclass(frozen=True, slots=True) @@ -15,28 +15,133 @@ class Complete: value: object +@dataclass(frozen=True, slots=True) +class Open: + value: None + + +@dataclass(frozen=True, slots=True) +class Yield: + value: object + + +Settled = Complete | Open | Yield +Step = Await | Settled + + class Execution(Protocol): - def start(self) -> Await | Complete: ... + def start(self) -> Step: ... - def resume_value(self, value: object) -> Await | Complete: ... + def resume_value(self, value: object) -> Step: ... - def resume_error(self, error: BaseException) -> Await | Complete: ... + def resume_error(self, error: BaseException) -> Step: ... def close(self) -> None: ... +class StreamClosed(Exception): + """Tells a streaming execution that its caller stopped reading.""" + + +async def _settle(execution: Execution, step: Step) -> Settled: + while isinstance(step, Await): + try: + value = await step.awaitable # rebind-ok: each selected await produces the next protocol input + except GeneratorExit: + raise + except BaseException as error: + step = execution.resume_error(error) # rebind-ok: advance the execution protocol + else: + step = execution.resume_value(value) # rebind-ok: advance the execution protocol + return step + + +def _settled(step: Step) -> Settled: + if isinstance(step, Await): + raise RuntimeError("sync call suspended") + return step + + async def drive(execution: Execution) -> object: + handed_off = False # rebind-ok: set once the execution belongs to the returned stream try: - step = execution.start() # rebind-ok: the execution protocol advances after each selected await - while isinstance(step, Await): - try: - value = await step.awaitable # rebind-ok: each selected await produces the next protocol input - except GeneratorExit: - raise - except BaseException as error: - step = execution.resume_error(error) # rebind-ok: advance the execution protocol - else: - step = execution.resume_value(value) # rebind-ok: advance the execution protocol + step: Final = await _settle(execution, execution.start()) + if isinstance(step, Open): + handed_off = True + return Stream(execution) return step.value finally: - execution.close() + if not handed_off: + execution.close() + + +class Stream(AsyncIterator[object]): + """A streamed native call: each read resumes the execution until its next chunk.""" + + def __init__(self, execution: Execution) -> None: + self._execution: Final = execution + self._done = False + + def __aiter__(self) -> Stream: + return self + + async def __anext__(self) -> object: + if self._done: + raise StopAsyncIteration + try: + step: Final = await _settle(self._execution, self._execution.resume_value(None)) + except BaseException: + self._finish() + raise + if isinstance(step, Yield): + return step.value + self._finish() + raise StopAsyncIteration + + async def aclose(self) -> None: + if self._done: + return + try: + await _settle(self._execution, self._execution.resume_error(StreamClosed())) + finally: + self._finish() + + def _finish(self) -> None: + self._done = True + self._execution.close() + + +class SyncStream(Iterator[object]): + """The sync form of `Stream`; its execution never suspends on an awaitable.""" + + def __init__(self, execution: Execution) -> None: + self._execution: Final = execution + self._done = False + + def __iter__(self) -> SyncStream: + return self + + def __next__(self) -> object: + if self._done: + raise StopIteration + try: + step: Final = _settled(self._execution.resume_value(None)) + except BaseException: + self._finish() + raise + if isinstance(step, Yield): + return step.value + self._finish() + raise StopIteration + + def close(self) -> None: + if self._done: + return + try: + _settled(self._execution.resume_error(StreamClosed())) + finally: + self._finish() + + def _finish(self) -> None: + self._done = True + self._execution.close() diff --git a/litellm/rust_bridge/messages/entrypoints.py b/litellm/rust_bridge/messages/entrypoints.py index 46565bfd46a..d25c906c4c1 100644 --- a/litellm/rust_bridge/messages/entrypoints.py +++ b/litellm/rust_bridge/messages/entrypoints.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping, Sequence from dataclasses import dataclass from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables @@ -26,7 +26,7 @@ class NativeMessages(Protocol): request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object], - ) -> AnthropicMessagesResponse: ... + ) -> AnthropicMessagesResponse | Iterator[bytes]: ... class NativeAmessages(Protocol): @@ -35,7 +35,7 @@ class NativeAmessages(Protocol): request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object], - ) -> Awaitable[AnthropicMessagesResponse]: ... + ) -> Awaitable[AnthropicMessagesResponse | AsyncIterator[bytes]]: ... def _messages_binding(value: object) -> NativeMessages | None: @@ -50,5 +50,5 @@ def _amessages_binding(value: object) -> NativeAmessages | None: return cast("NativeAmessages", value) # cast-ok: callable validated at the native binding boundary -NATIVE_MESSAGES: Final = NativeBinding("anthropic_messages_handler", validate=_messages_binding) -NATIVE_AMESSAGES: Final = NativeBinding("anthropic_messages", validate=_amessages_binding) +NATIVE_MESSAGES: Final = NativeBinding("messages", validate=_messages_binding) +NATIVE_AMESSAGES: Final = NativeBinding("amessages", validate=_amessages_binding) diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/messages/route_host.py index 1aff6c7f75d..beef0f81eca 100644 --- a/litellm/rust_bridge/messages/route_host.py +++ b/litellm/rust_bridge/messages/route_host.py @@ -20,4 +20,4 @@ def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]: def map_failure(error: Exception, request: LiteLLMMessagesRequest, request_provider: str) -> Exception: - return failures.map_failure(error, request.model, request_provider, arguments(request)) + return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base) diff --git a/litellm/rust_bridge/ocr/route_host.py b/litellm/rust_bridge/ocr/route_host.py index 277fdceb734..bfbd5c11d4e 100644 --- a/litellm/rust_bridge/ocr/route_host.py +++ b/litellm/rust_bridge/ocr/route_host.py @@ -4,40 +4,17 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final -import httpx -import openai -from pydantic import TypeAdapter, ValidationError +from pydantic import TypeAdapter import litellm from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse from litellm.rust_bridge import failures +from litellm.rust_bridge.failures import UpstreamFailure from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest +__all__ = ("UpstreamFailure", "arguments", "map_failure", "response") + _RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object]) -_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) -_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) - - -class UpstreamFailure(Exception): - def __init__(self, response: httpx.Response, cause: Exception) -> None: - super().__init__(str(cause)) - self.message: Final = str(cause) - self.response: Final = response - self.status_code: Final = response.status_code - self.__cause__ = cause - - -def _upstream_failure(error: Exception, request: LiteLLMOcrRequest) -> Exception: - try: - status, body = _UPSTREAM_ARGS.validate_python(error.args) - headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) - except ValidationError: - return error - http_request: Final = httpx.Request("POST", request.api_base or "https://docs.litellm.ai/docs") - return UpstreamFailure( - httpx.Response(status, content=body.encode(), headers=headers, request=http_request), - error, - ) def response(value: Mapping[str, object]) -> OCRResponse: @@ -61,11 +38,4 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: model=request.model.removeprefix(f"{request_provider}/"), llm_provider=request_provider, ) - original: Final = _upstream_failure(error, request) - public_error: Final = failures.map_failure(original, request.model, request_provider, arguments(request)) - if isinstance(original, UpstreamFailure) and public_error.__context__ is original: - public_error.__context__ = error - if isinstance(public_error, openai.APIStatusError): - public_error.response = original.response - public_error.status_code = original.status_code - return public_error + return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base) diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 4fa4c0b95ec..0b442f1f269 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -73,7 +73,7 @@ def assert_native_request( headers: HTTPMessage, body: object, ) -> None: - if route not in {"transcription", "messages", "chat_completions"}: + if route not in {"transcription", "chat_completions"}: raise AssertionError(f"unexpected route marker: {route!r}") if outcome not in {"success", "429", "hang"}: raise AssertionError(f"unexpected outcome marker: {outcome!r}") @@ -89,10 +89,6 @@ def assert_native_request( assert path == "/v1/messages" assert headers.get("x-api-key") == "sk-native" assert body["model"] == "claude-sonnet-4-5" - if route == "messages": - assert body["max_tokens"] == 16 - assert body["messages"][0]["content"] == "hello-from-messages" - return assert body["max_tokens"] == 17 assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}] @@ -132,17 +128,6 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: "language": "en", }, } - if route == "messages": - return common | { - "model": "claude-sonnet-4-5", - "body": { - "model": "claude-sonnet-4-5", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hello-from-messages"}], - }, - "api_key": "sk-native", - "custom_llm_provider": "anthropic", - } if route == "chat_completions": return common | { "model": "anthropic/claude-sonnet-4-5", @@ -165,8 +150,6 @@ def assert_success(route: str, response: object) -> None: def success_value(route: str, response: dict[object, object]) -> object: if route == "transcription": return response["text"] - if route == "messages": - return response["content"][0]["text"] return response["choices"][0]["message"]["content"] @@ -181,7 +164,7 @@ def assert_rate_limit(native: object, route: str, error: BaseException) -> None: def exercise_sync(native: object, api_base: str) -> None: - for route in ("transcription", "messages", "chat_completions"): + for route in ("transcription", "chat_completions"): function: Final = getattr(native, route) assert_success(route, function(**route_kwargs(route, api_base, "success"))) try: @@ -193,7 +176,7 @@ def exercise_sync(native: object, api_base: str) -> None: async def exercise_async(native: object, api_base: str) -> None: - for route in ("transcription", "messages", "chat_completions"): + for route in ("transcription", "chat_completions"): function: Final = getattr(native, f"a{route}") assert_success(route, await function(**route_kwargs(route, api_base, "success"))) try: @@ -206,11 +189,11 @@ async def exercise_async(native: object, api_base: str) -> None: async def exercise_async_concurrency(native: object, api_base: str) -> None: responses: Final = await asyncio.wait_for( - asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))), + asyncio.gather(*(native.achat_completions(**route_kwargs("chat_completions", api_base, "success")) for _ in range(32))), timeout=15, ) for response in responses: - assert_success("messages", response) + assert_success("chat_completions", response) def exercise_routes(native_path: Path, api_base: str) -> object: @@ -223,8 +206,8 @@ def exercise_routes(native_path: Path, api_base: str) -> object: def exercise_signal(native: object, api_base: str) -> int: try: - native.messages( - **route_kwargs("messages", api_base, "hang"), + native.chat_completions( + **route_kwargs("chat_completions", api_base, "hang"), ) except KeyboardInterrupt: sys.stdout.write("KeyboardInterrupt\n") diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py index 72390b79141..b882a1bb8c2 100644 --- a/tests/test_litellm/rust_bridge/test_bindings.py +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -43,8 +43,8 @@ def test_binding_validates_native_attribute( ROUTE_BINDINGS: Final = ( ("completion", chat_completions.NATIVE_COMPLETION), ("acompletion", chat_completions.NATIVE_ACOMPLETION), - ("anthropic_messages_handler", messages.NATIVE_MESSAGES), - ("anthropic_messages", messages.NATIVE_AMESSAGES), + ("messages", messages.NATIVE_MESSAGES), + ("amessages", messages.NATIVE_AMESSAGES), ("responses", responses.NATIVE_RESPONSES), ("aresponses", responses.NATIVE_ARESPONSES), ("ocr", ocr.NATIVE_OCR), diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index 2c737b0160e..e9fdbf859f4 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -40,6 +40,10 @@ def test_shipped_decisions( enabled: Final = environment == "1" if environment is not None else process is not False assert catalog.rollout(context) is Rollout.RUST_OPT_OUT assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) + elif route is Route.MESSAGES: + enabled: Final = environment == "1" if environment is not None else process is True + assert catalog.rollout(context) is Rollout.RUST_OPT_IN + assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) elif route is Route.TRANSCRIPTION and provider == "bedrock": assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index ade0ae549fb..fa6c0b30413 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -157,7 +157,6 @@ def test_context_outside_rule_stays_on_python() -> None: ( Context(Route.CHAT_COMPLETIONS, provider="anthropic"), Context(Route.CHAT_COMPLETIONS, provider="bedrock"), - Context(Route.MESSAGES, provider="anthropic"), Context(Route.RESPONSES, provider="openai"), Context(Route.TRANSCRIPTION, provider="openai"), ), diff --git a/tests/test_litellm_rust/messages/__init__.py b/tests/test_litellm_rust/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py new file mode 100644 index 00000000000..b55bc47d640 --- /dev/null +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -0,0 +1,175 @@ +from collections.abc import AsyncIterator, Iterator +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import ( + MESSAGES, + MESSAGES_EVENTS, + MESSAGES_MODEL, + MESSAGES_RESPONSE, + request_body, +) + +pytestmark = pytest.mark.requires_rust_extension + +STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS) + + +@pytest.fixture +def messages_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server + + +def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: + return { + "model": MESSAGES_MODEL, + "messages": [dict(message) for message in MESSAGES], + "max_tokens": 64, + "api_key": "test-key", + "api_base": server.base_url, + **kwargs, + } + + +def assert_served_natively(server: RecordingServer) -> None: + assert len(server.requests) == 1 + assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx") + + +@pytest.mark.asyncio +async def test_native_messages_callbacks_see_the_provider_request_and_the_public_response( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + response: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, callbacks=[recorder], litellm_call_id="messages-success") + ) + + assert_served_natively(messages_server) + assert response["content"] == MESSAGES_RESPONSE["content"] + sent: Final = messages_server.requests[0] + assert sent.path == "/v1/messages" + assert sent.body == {"model": "claude-sonnet-5", "messages": list(MESSAGES), "max_tokens": 64, "stream": False} + pre_call: Final = recorder.wait_for("log_pre_api_call") + assert request_body(pre_call[0].kwargs) == sent.body + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].call_type == "anthropic_messages" + assert success[0].kwargs["litellm_call_id"] == "messages-success" + assert success[0].response.choices[0].message.content == "Hello from native Messages" + + +@pytest.mark.asyncio +async def test_native_messages_pre_call_body_edit_reaches_the_provider(messages_server: RecordingServer) -> None: + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs)["temperature"] = 0.25 + + await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Edit()])) + + assert messages_server.requests[0].body["temperature"] == 0.25 + + +@pytest.mark.asyncio +async def test_native_messages_provider_error_reaches_caller_and_failure_callbacks_as_one_public_error( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue( + ResponseSpec(body={"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}}, status=400) + ) + observed: Final = [] + + class Observe(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs["exception"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs["exception"])) + + with pytest.raises(litellm.BadRequestError) as raised: + await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Observe()])) + + assert_served_natively(messages_server) + assert [phase for phase, _ in observed] == ["sync", "async"] + assert all(error is raised.value for _, error in observed) + + +def sse_payload() -> bytes: + return b"".join(STREAM.payloads()) + + +@pytest.mark.asyncio +async def test_native_messages_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, stream=True, callbacks=[recorder]) + ) + assert isinstance(stream, AsyncIterator) + first: Final = await anext(stream) + await drain_logging() + assert "async_log_success_event" not in recorder.names + rest: Final = [chunk async for chunk in stream] + + assert first + b"".join(rest) == sse_payload() + assert_served_natively(messages_server) + assert messages_server.requests[0].body["stream"] is True + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].kwargs["stream"] is True + assert success[0].kwargs["completion_start_time"] is not None + assert "log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_native_messages_stream_closed_early_logs_success_once_for_what_was_delivered( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, stream=True, callbacks=[recorder]) + ) + assert isinstance(stream, AsyncIterator) + await anext(stream) + await stream.aclose() + + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + with pytest.raises(StopAsyncIteration): + await anext(stream) + + +def test_native_sync_messages_stream_relays_provider_events_and_logs_success_once( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) + assert isinstance(stream, Iterator) + + assert b"".join(stream) == sse_payload() + assert_served_natively(messages_server) + assert len(recorder.wait_for("async_log_success_event")) == 1 + + +def test_native_sync_messages_returns_the_provider_message(messages_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + response: Final = litellm.anthropic.messages.create(**arguments(messages_server, callbacks=[recorder])) + + assert_served_natively(messages_server) + assert response["content"] == MESSAGES_RESPONSE["content"] + assert len(recorder.wait_for("log_success_event")) == 1 diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index 228ed2cc454..3eea47751d3 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -25,6 +25,12 @@ class ResponseSpec: status: int = 200 headers: dict[str, str] = field(default_factory=dict) delay: float = 0 + events: tuple[tuple[str, object], ...] = () + + def payloads(self) -> tuple[bytes, ...]: + if not self.events: + return (json.dumps(self.body).encode(),) + return tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in self.events) @dataclass @@ -73,15 +79,17 @@ def recording_service() -> Iterator[RecordingServer]: response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response) if response.delay: time.sleep(response.delay) - payload: Final = json.dumps(response.body).encode() + payloads: Final = response.payloads() self.send_response(response.status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(payload))) + self.send_header("Content-Type", "text/event-stream" if response.events else "application/json") + self.send_header("Content-Length", str(sum(len(payload) for payload in payloads))) for name, value in response.headers.items(): self.send_header(name, value) self.end_headers() try: - self.wfile.write(payload) + for payload in payloads: + self.wfile.write(payload) + self.wfile.flush() except (BrokenPipeError, ConnectionResetError): pass diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py index b60cf5eac02..c9cf81b83ca 100644 --- a/tests/test_litellm_rust/support/requests.py +++ b/tests/test_litellm_rust/support/requests.py @@ -12,6 +12,41 @@ OCR_RESPONSE: Final = { "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, } +MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5" +MESSAGES: Final = ({"role": "user", "content": "Hello"},) +MESSAGES_RESPONSE: Final = { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "Hello from native Messages"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 4}, +} +MESSAGES_EVENTS: Final = ( + ("message_start", {"type": "message_start", "message": {**MESSAGES_RESPONSE, "content": [], "stop_reason": None}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello from native Messages"}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + }, + ), + ("message_stop", {"type": "message_stop"}), +) + def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: return { From 3a5b7c12ef119474a596cd58f59807cce5854fb5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:35:20 -0700 Subject: [PATCH 094/102] refactor(rust): separate machine events from the Python lifecycle's events --- .../crates/callbacks-legacy/src/adapter.rs | 57 +++++++-------- .../tests/deployment_hooks.rs | 11 ++- .../crates/callbacks-legacy/tests/payload.rs | 8 +-- .../crates/callbacks-legacy/tests/terminal.rs | 14 ++-- litellm-rust/crates/callbacks/src/event.rs | 16 +++-- litellm-rust/crates/callbacks/src/host.rs | 4 +- litellm-rust/crates/callbacks/src/run.rs | 5 +- litellm-rust/crates/core/src/machine/mod.rs | 4 +- .../crates/core/src/messages/route.rs | 4 +- litellm-rust/crates/core/src/ocr/handler.rs | 4 +- .../tests/azure_document_intelligence_ocr.rs | 8 +-- litellm-rust/crates/core/tests/ocr.rs | 9 ++- litellm-rust/crates/core/tests/reducto_ocr.rs | 8 +-- .../crates/host-python/src/adapter.rs | 33 ++++++--- litellm-rust/crates/host-python/src/driver.rs | 69 ++++++++++--------- litellm-rust/crates/host-python/src/lib.rs | 2 +- 16 files changed, 142 insertions(+), 114 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 9d28db92add..a67da2188fb 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -3,10 +3,10 @@ //! `@client` path makes them. use litellm_callbacks::event::{ - CallEvent, FailureOrigin, RequestContext, Timing, WireRequest, epoch_seconds, + FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds, }; use litellm_host_python::{ - LifecycleStep, PublicValue, PythonLifecycle, from_py, missing_state, to_py, + LifecycleEvent, LifecycleStep, PythonLifecycle, from_py, missing_state, to_py, }; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -346,28 +346,11 @@ impl PythonLifecycle for LegacyLogging { fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult { - match (event, public) { - (CallEvent::Started { .. }, _) => Ok(LifecycleStep::Done), - (CallEvent::Opened, _) => { - Streaming::Opened.call(py, (self.logger()?.object(py),))?; - self.stream = Some(DeliveredStream { - chunks: PyList::empty(py).unbind(), - first_chunk: None, - }); - Ok(LifecycleStep::Done) - } - (CallEvent::Delivered, Some(PublicValue::Chunk(chunk))) => { - let stream = self.stream.as_mut().ok_or_else(missing_state)?; - if stream.first_chunk.is_none() { - stream.first_chunk = Some(datetime(py, epoch_seconds())?); - } - stream.chunks.bind(py).append(chunk)?; - Ok(LifecycleStep::Done) - } - (CallEvent::ResponseReceived { raw }, _) => { + match event { + LifecycleEvent::Started { .. } => Ok(LifecycleStep::Done), + LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { let api_key = self .context .as_ref() @@ -382,7 +365,7 @@ impl PythonLifecycle for LegacyLogging { )?; Ok(LifecycleStep::Done) } - (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { + LifecycleEvent::Succeeded { timing, response } => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); match &self.stream { @@ -391,13 +374,17 @@ impl PythonLifecycle for LegacyLogging { } Ok(LifecycleStep::Done) } - (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { + LifecycleEvent::Failed { + timing, + origin, + error, + } => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); if self.stream.is_some() { return self.stream_failure(py); } - if *origin == FailureOrigin::Call + if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() { @@ -412,10 +399,26 @@ impl PythonLifecycle for LegacyLogging { } self.dispatch_failure(py) } - _ => Err(missing_state()), } } + fn opened(&mut self, py: Python<'_>) -> PyResult<()> { + Streaming::Opened.call(py, (self.logger()?.object(py),))?; + self.stream = Some(DeliveredStream { + chunks: PyList::empty(py).unbind(), + first_chunk: None, + }); + Ok(()) + } + + fn delivered(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()> { + let stream = self.stream.as_mut().ok_or_else(missing_state)?; + if stream.first_chunk.is_none() { + stream.first_chunk = Some(datetime(py, epoch_seconds())?); + } + stream.chunks.bind(py).append(chunk) + } + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { match self.pending.take().ok_or_else(missing_state)? { Pending::DeploymentPreCall => { diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index 7bad09c7890..f4602739548 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; +use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -217,13 +217,12 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle .resume(py, Ok(local(&locals, "kwargs").unbind())) .unwrap(); let failure = PyErr::from_value(local(&locals, "failure")); - let failed = CallEvent::Failed { + let failed = LifecycleEvent::Failed { timing: TIMING, origin: FailureOrigin::Call, + error: &failure, }; - let step = logging - .emit(py, &failed, Some(PublicValue::Error(&failure))) - .unwrap(); + let step = logging.emit(py, failed).unwrap(); assert!(awaits_deployment_hook(&step)); let hook_result = if cancelled { Err(CancelledError::new_err("cancelled")) diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 43128bc38ea..67ad4ab8a2a 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,8 +1,8 @@ use std::ffi::CStr; use litellm_auth::SecretValue; -use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{LifecycleStep, PythonLifecycle, to_py}; +use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; use proptest::prelude::*; use pyo3::prelude::*; use rstest::rstest; @@ -94,13 +94,13 @@ fn before_send_bound( body, }; let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = CallEvent::ResponseReceived { + let raw = MachineEvent::ResponseReceived { raw: RawResponse { body: "raw response".into(), }, }; assert!(matches!( - logging.emit(py, &raw, None).unwrap(), + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), LifecycleStep::Done )); run(py, &locals, c"check()"); diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 3094d7b88d2..5688f70f387 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; +use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; @@ -33,8 +33,10 @@ fn succeed( logging .emit( py, - &CallEvent::Succeeded { timing: TIMING }, - Some(PublicValue::Response(&response)), + LifecycleEvent::Succeeded { + timing: TIMING, + response: &response, + }, ) .unwrap() } @@ -44,11 +46,11 @@ fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) logging .emit( py, - &CallEvent::Failed { + LifecycleEvent::Failed { timing: TIMING, origin: FailureOrigin::Host, + error: &failure, }, - Some(PublicValue::Error(&failure)), ) .unwrap() } diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs index d19e973b812..182dab657d3 100644 --- a/litellm-rust/crates/callbacks/src/event.rs +++ b/litellm-rust/crates/callbacks/src/event.rs @@ -52,18 +52,20 @@ pub enum FailureOrigin { Host, } +/// What a machine reports while it runs. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum MachineEvent { + ResponseReceived { raw: RawResponse }, +} + +/// What an in-process host observes: the machine's own events between the driver's +/// start and terminal ones. #[derive(Clone, Debug, PartialEq)] pub enum CallEvent { Started { start_time: f64, }, - ResponseReceived { - raw: RawResponse, - }, - /// The call streams and its stream was handed to the caller. - Opened, - /// One chunk of an open stream reached the caller. - Delivered, + Machine(MachineEvent), Succeeded { timing: Timing, }, diff --git a/litellm-rust/crates/callbacks/src/host.rs b/litellm-rust/crates/callbacks/src/host.rs index eef3e1da8d5..aba35185a18 100644 --- a/litellm-rust/crates/callbacks/src/host.rs +++ b/litellm-rust/crates/callbacks/src/host.rs @@ -1,6 +1,6 @@ use std::future::Future; -use crate::event::{CallEvent, RequestContext, WireRequest}; +use crate::event::{CallEvent, MachineEvent, RequestContext, WireRequest}; use crate::route::Route; /// One suspension point of a native call, performed by the host. @@ -10,7 +10,7 @@ pub enum HostOp { wire: Box, context: Box, }, - Emit(CallEvent), + Emit(MachineEvent), /// The response streams: the host hands the caller a stream and answers once the /// caller asks for the first chunk or goes away. Open(R::StreamHead), diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/callbacks/src/run.rs index 705c5504c01..6a0c08fba68 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/callbacks/src/run.rs @@ -25,7 +25,10 @@ where .before_send(*wire, &context) .await .map(|wire| HostResult::BeforeSend(Box::new(wire))), - HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted), + HostOp::Emit(event) => host + .emit(&CallEvent::Machine(event)) + .await + .map(|()| HostResult::Emitted), HostOp::Open(head) => host.open(head).await.map(HostResult::Demand), HostOp::Deliver(chunk) => host.deliver(chunk).await.map(HostResult::Demand), }; diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/core/src/machine/mod.rs index 929f0a423c4..d6db488159e 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/core/src/machine/mod.rs @@ -8,7 +8,7 @@ use std::{future::Future, pin::Pin}; pub use auth::{HostTokenProvider, TokenRoute}; use litellm_callbacks::{ - event::{CallEvent, RequestContext, WireRequest}, + event::{MachineEvent, RequestContext, WireRequest}, host::{Demand, HostOp, HostResult}, machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, route::Route, @@ -82,7 +82,7 @@ where } } - pub async fn emit(&self, event: CallEvent) -> Result<(), R::Error> { + pub async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> { match self.invoke(HostOp::Emit(event)).await? { HostResult::Emitted => Ok(()), _ => Err(MachineFault::Mismatch.into()), diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index a680b5ce1ec..2fd2f79907a 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -3,7 +3,7 @@ use std::{sync::Mutex, time::Duration}; use bytes::Bytes; use litellm_auth::SecretValue; use litellm_callbacks::{ - event::{CallEvent, RawResponse, RequestContext, WireRequest}, + event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, route::Route, }; @@ -172,7 +172,7 @@ async fn execute(host: MessagesHost) -> Result { return relay(&host, response).await; } let text = response.text().await.map_err(network)?; - host.emit(CallEvent::ResponseReceived { + host.emit(MachineEvent::ResponseReceived { raw: RawResponse { body: text.clone() }, }) .await?; diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index a6af190eb91..aff1eded2cc 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,6 +1,6 @@ use futures_util::future::BoxFuture; use litellm_auth::SecretValue; -use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; +use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, @@ -71,7 +71,7 @@ impl CallHooks for OcrCallHooks { } fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(self.host.emit(CallEvent::ResponseReceived { + Box::pin(self.host.emit(MachineEvent::ResponseReceived { raw: RawResponse { body: String::from_utf8_lossy(body).into_owned(), }, diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 1fc4d6c2b9e..62544931dfc 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::CallEvent; +use litellm_callbacks::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::error::Error; use rstest::rstest; use serde_json::{Value, json}; @@ -263,7 +263,7 @@ async fn accepted_response_emits_response_received_before_polling() { json!({}), )) .with_observer(move |event| { - let CallEvent::ResponseReceived { raw } = event else { + let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { return; }; match request_count.lock().unwrap().len() { @@ -466,7 +466,7 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() { mod transformation { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::CallEvent; + use litellm_callbacks::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::transformation::OcrDocument; use serde_json::{Value, json}; @@ -646,7 +646,7 @@ mod transformation { json!({}), )) .with_observer(move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { observed .lock() .unwrap() diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 779d2037bb3..a6001951361 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use litellm_callbacks::{ - event::{CallEvent, WireRequest}, + event::{CallEvent, MachineEvent, WireRequest}, host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; @@ -195,11 +195,9 @@ async fn facade_uses_the_injected_http_client() { fn event_name(event: &CallEvent) -> &'static str { match event { CallEvent::Started { .. } => "started", - CallEvent::ResponseReceived { .. } => "response", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", - CallEvent::Opened => "opened", - CallEvent::Delivered => "delivered", } } @@ -377,6 +375,7 @@ async fn drive_until( intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) } HostOp::Emit(event) => { + let event = CallEvent::Machine(event); ops.push(event_name(&event)); host.emit(&event) .await @@ -420,7 +419,7 @@ async fn invalid_provider_response_emits_response_received_before_normalization_ let observed = responses_received.clone(); let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_observer( move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { observed.lock().unwrap().push(raw.body.clone()); } }, diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 59891b16e90..8c037889a17 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, WireRequest}; +use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; use rstest::rstest; use serde_json::{Value, json}; @@ -139,7 +139,7 @@ async fn response_received_stays_after_reducto_upload_and_parse() { let request_count = seen.clone(); let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer( move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { assert_eq!(request_count.lock().unwrap().len(), 2); assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); } @@ -351,7 +351,7 @@ async fn guardrail_rewrites_document_before_upload() { } mod transformation { - use litellm_callbacks::event::{CallEvent, WireRequest}; + use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::{ base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, reducto::ocr::transformation::*, @@ -506,7 +506,7 @@ mod transformation { let request_count = seen.clone(); let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) .with_observer(move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { assert_eq!(request_count.lock().unwrap().len(), 2); assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); } diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 795977438fd..f946dbc763a 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, RequestContext, Timing, WireRequest}; +use litellm_callbacks::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; use litellm_callbacks::route::Route; use pyo3::exceptions::PyRuntimeError; use pyo3::gc::{PyTraverseError, PyVisit}; @@ -19,11 +19,22 @@ pub enum LifecycleStep { Done, } -/// The host-typed value the driver attaches to a terminal event. -pub enum PublicValue<'a> { - Response(&'a Py), - Error(&'a PyErr), - Chunk(&'a Py), +/// What a lifecycle observes: the driver's start, the machine's own events, and one +/// terminal event carrying the public value the caller receives. +pub enum LifecycleEvent<'a> { + Started { + start_time: f64, + }, + Machine(&'a MachineEvent), + Succeeded { + timing: Timing, + response: &'a Py, + }, + Failed { + timing: Timing, + origin: FailureOrigin, + error: &'a PyErr, + }, } /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in @@ -61,10 +72,16 @@ pub trait PythonLifecycle: Send + Sync { fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult; + /// The call streams and its stream was handed to the caller. The caller is not + /// inside an await here, so this step and `delivered` cannot suspend. + fn opened(&mut self, py: Python<'_>) -> PyResult<()>; + + /// One chunk of an open stream is about to reach the caller. + fn delivered(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()>; + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; fn close(&mut self, py: Python<'_>); diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index c59ba1925dc..25f78d7013e 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing, epoch_seconds}; +use litellm_callbacks::event::{FailureOrigin, Timing, epoch_seconds}; use litellm_callbacks::host::{Demand, HostOp, HostResult, HostStep}; use litellm_callbacks::machine::{HostFailure, Machine, MachineStep}; use litellm_callbacks::route::Route; @@ -13,7 +13,7 @@ use pyo3::types::PyDict; use tokio::sync::Mutex; use crate::adapter::{ - HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, + HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -159,10 +159,10 @@ where match (self.pending.take(), result) { (None, None) => { self.started_at = epoch_seconds(); - let started = CallEvent::Started { + let started = LifecycleEvent::Started { start_time: self.started_at, }; - match self.adapter.emit(py, &started, None) { + match self.adapter.emit(py, started) { Ok(step) => self.on_adapter(py, step, Expect::Started), Err(error) => self.adapter_failed(py, error), } @@ -303,7 +303,7 @@ where } HostOp::Open(_) => return self.opened(py).map(Next::Return), HostOp::Deliver(chunk) => return self.delivered(py, chunk).map(Next::Return), - HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { + HostOp::Emit(event) => match self.adapter.emit(py, LifecycleEvent::Machine(&event)) { Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Emitted)); @@ -321,12 +321,11 @@ where fn opened(&mut self, py: Python<'_>) -> PyResult { self.stage = Stage::Streaming; - match self.adapter.emit(py, &CallEvent::Opened, None) { - Ok(LifecycleStep::Done) => { + match self.adapter.opened(py) { + Ok(()) => { self.pending = Some(Pending::Consumer); Ok(ExecutionStep::Open) } - Ok(_) => Err(missing_state()), Err(error) => self.interrupt(py, error), } } @@ -340,15 +339,11 @@ where Ok(chunk) => chunk, Err(error) => return self.interrupt(py, error), }; - let observed = - self.adapter - .emit(py, &CallEvent::Delivered, Some(PublicValue::Chunk(&chunk))); - match observed { - Ok(LifecycleStep::Done) => { + match self.adapter.delivered(py, &chunk) { + Ok(()) => { self.pending = Some(Pending::Consumer); Ok(ExecutionStep::Yield(chunk)) } - Ok(_) => Err(missing_state()), Err(error) => self.interrupt(py, error), } } @@ -461,12 +456,11 @@ where } fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { - let event = CallEvent::Succeeded { + let event = LifecycleEvent::Succeeded { timing: self.timing(), + response: &response, }; - let step = self - .adapter - .emit(py, &event, Some(PublicValue::Response(&response)))?; + let step = self.adapter.emit(py, event)?; self.stage = Stage::Succeeded(response); self.on_adapter(py, step, Expect::Terminal) } @@ -481,13 +475,12 @@ where if is_cancellation(py, &error) { return Err(error); } - let event = CallEvent::Failed { + let event = LifecycleEvent::Failed { timing: self.timing(), origin, + error: &error, }; - let step = self - .adapter - .emit(py, &event, Some(PublicValue::Error(&error)))?; + let step = self.adapter.emit(py, event)?; self.stage = Stage::Failed(error.into_value(py)); self.on_adapter(py, step, Expect::Terminal) } @@ -541,7 +534,7 @@ where mod tests { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::{RequestContext, WireRequest}; + use litellm_callbacks::event::{MachineEvent, RequestContext, WireRequest}; use litellm_callbacks::machine::{Interrupted, Step}; use pyo3::exceptions::{PyBaseException, PyValueError}; use pyo3::types::PyDict; @@ -803,23 +796,33 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult { - self.log.push(match (event, public) { - (CallEvent::Started { .. }, None) => "started".into(), - (CallEvent::ResponseReceived { raw }, None) => format!("response:{}", raw.body), - (CallEvent::Succeeded { .. }, Some(PublicValue::Response(value))) => { - format!("succeeded:{}", value.bind(py)) + self.log.push(match event { + LifecycleEvent::Started { .. } => "started".into(), + LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { + format!("response:{}", raw.body) } - (CallEvent::Failed { origin, .. }, Some(PublicValue::Error(error))) => { + LifecycleEvent::Succeeded { response, .. } => { + format!("succeeded:{}", response.bind(py)) + } + LifecycleEvent::Failed { origin, error, .. } => { format!("failed:{origin:?}:{}", error.value(py)) } - _ => "unexpected".into(), }); Ok(LifecycleStep::Done) } + fn opened(&mut self, _: Python<'_>) -> PyResult<()> { + self.log.push("opened"); + Ok(()) + } + + fn delivered(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { + self.log.push("delivered"); + Ok(()) + } + fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { Err(missing_state()) } @@ -899,7 +902,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri wire: Box::new(wire()), context: Box::new(context()), }, - HostOp::Emit(CallEvent::ResponseReceived { + HostOp::Emit(MachineEvent::ResponseReceived { raw: litellm_callbacks::event::RawResponse { body: "raw".into() }, }), ], diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8889b1513db..3738a9b1c3c 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -13,7 +13,7 @@ mod handle; mod marshal; pub use adapter::{ - HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, + HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; pub use argument::lookup; pub use callable::wrap_failure; From c2679757b3ad93bd9f8def74021e78be42a025c7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:36:42 -0700 Subject: [PATCH 095/102] refactor(rust): put stream billing on the legacy surface --- .../callbacks-legacy/python_contract.json | 3 ++ .../crates/callbacks-legacy/src/adapter.rs | 37 +++++++++++++++---- .../crates/callbacks-legacy/src/lib.rs | 2 +- .../crates/callbacks-legacy/tests/support.rs | 1 + .../crates/host-python/src/adapter.rs | 6 +-- litellm-rust/crates/host-python/src/driver.rs | 6 +-- .../python-bridge/src/routes/messages/mod.rs | 6 ++- .../python-bridge/src/routes/ocr/mod.rs | 1 + litellm/rust_bridge/legacy_callbacks.py | 14 +++++-- 9 files changed, 53 insertions(+), 23 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy/python_contract.json index a09bdc711a3..8a7f3b98f47 100644 --- a/litellm-rust/crates/callbacks-legacy/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy/python_contract.json @@ -100,6 +100,8 @@ ], "stream_success": [ "logger", + "url_route", + "endpoint_type", "request_body", "chunks", "start", @@ -108,6 +110,7 @@ ], "stream_failure": [ "logger", + "endpoint_type", "request_body", "chunks", "error" diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index a67da2188fb..4e9167100e0 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -30,6 +30,16 @@ pub struct LegacySurface { pub call_type: &'static str, /// What `Logging.pre_call` is told the input was. pub input_description: &'static str, + /// How a streamed response is billed; `None` for a route that never streams. + pub stream: Option, +} + +/// The pass-through billing a streamed response goes through once its chunks are in. +#[derive(Clone, Copy, Debug)] +pub struct PassThroughStream { + pub url_route: &'static str, + /// A value of Python's `EndpointType`. + pub endpoint_type: &'static str, } /// What the Messages stream iterator keeps for its end-of-stream billing. @@ -177,10 +187,13 @@ impl LegacyLogging { fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> { let logger = self.logger()?; + let billing = self.surface.stream.ok_or_else(missing_state)?; let billed = Streaming::Success.call( py, ( logger.object(py), + billing.url_route, + billing.endpoint_type, &self.body, &stream.chunks, &self.start, @@ -201,14 +214,25 @@ impl LegacyLogging { /// partial usage. The sync path has no loop to schedule that on, so it falls back to /// the plain failure handler. fn stream_failure(&mut self, py: Python<'_>) -> PyResult { - let (Some(logger), Some(error), Some(stream)) = (&self.logger, &self.error, &self.stream) + let (Some(logger), Some(error), Some(stream), Some(billing)) = + (&self.logger, &self.error, &self.stream, self.surface.stream) else { return Ok(LifecycleStep::Done); }; if !self.asynchronous { return self.dispatch_failure(py); } - match Streaming::Failure.call(py, (logger.object(py), &self.body, &stream.chunks, error)) { + let scheduled = Streaming::Failure.call( + py, + ( + logger.object(py), + billing.endpoint_type, + &self.body, + &stream.chunks, + error, + ), + ); + match scheduled { Ok(awaitable) => { self.pending = Some(Pending::AsyncFailure); Ok(LifecycleStep::Await(awaitable.unbind())) @@ -343,11 +367,7 @@ impl PythonLifecycle for LegacyLogging { self.finalize(py) } - fn emit( - &mut self, - py: Python<'_>, - event: LifecycleEvent<'_>, - ) -> PyResult { + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult { match event { LifecycleEvent::Started { .. } => Ok(LifecycleStep::Done), LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { @@ -403,6 +423,9 @@ impl PythonLifecycle for LegacyLogging { } fn opened(&mut self, py: Python<'_>) -> PyResult<()> { + if self.surface.stream.is_none() { + return Err(missing_state()); + } Streaming::Opened.call(py, (self.logger()?.object(py),))?; self.stream = Some(DeliveredStream { chunks: PyList::empty(py).unbind(), diff --git a/litellm-rust/crates/callbacks-legacy/src/lib.rs b/litellm-rust/crates/callbacks-legacy/src/lib.rs index 42ffd545e2b..eaa1a8b714e 100644 --- a/litellm-rust/crates/callbacks-legacy/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy/src/lib.rs @@ -21,7 +21,7 @@ mod preparation; mod test_support; pub(crate) use adapter::LegacyLogging; -pub use adapter::LegacySurface; +pub use adapter::{LegacySurface, PassThroughStream}; pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy/tests/support.rs index 42ca184eb16..d3cc32e301f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/support.rs @@ -197,6 +197,7 @@ pub(crate) fn legacy_call( LegacySurface { call_type: "test", input_description: "test input", + stream: None, }, call, asynchronous, diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index f946dbc763a..e70e4a8c58f 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -69,11 +69,7 @@ pub trait PythonLifecycle: Send + Sync { timing: Timing, ) -> PyResult; - fn emit( - &mut self, - py: Python<'_>, - event: LifecycleEvent<'_>, - ) -> PyResult; + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult; /// The call streams and its stream was handed to the caller. The caller is not /// inside an await here, so this step and `delivered` cannot suspend. diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 25f78d7013e..d6b31b27314 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -793,11 +793,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } - fn emit( - &mut self, - py: Python<'_>, - event: LifecycleEvent<'_>, - ) -> PyResult { + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult { self.log.push(match event { LifecycleEvent::Started { .. } => "started".into(), LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index b4259aca5ba..8c42315ac59 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,7 +1,7 @@ mod host; use host::MessagesRouteHost; -use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; +use litellm_callbacks_legacy::{LegacySurface, PassThroughStream, PublicCall, run_legacy_call}; use litellm_core::messages::route::{messages_machine, supports}; use pyo3::{ prelude::*, @@ -13,6 +13,10 @@ use crate::errors::RustBridgeDeclined; const SURFACE: LegacySurface = LegacySurface { call_type: "anthropic_messages", input_description: "Messages", + stream: Some(PassThroughStream { + url_route: "/v1/messages", + endpoint_type: "anthropic", + }), }; fn run_messages( diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index b5bb941708d..8afa1e2a906 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -15,6 +15,7 @@ use pyo3::{ const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", input_description: "OCR document processing", + stream: None, }; const ASYNC_SURFACE: LegacySurface = LegacySurface { diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index fd0fac1799a..bac40442ce5 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -331,6 +331,8 @@ def stream_opened(logger: Logging) -> None: def stream_success( logger: Logging, + url_route: str, + endpoint_type: str, request_body: dict[str, object], chunks: list[bytes], start: datetime.datetime, @@ -353,9 +355,9 @@ def stream_success( coroutine: Final = build( litellm_logging_obj=logger, passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, - url_route="/v1/messages", + url_route=url_route, request_body=request_body, - endpoint_type=EndpointType.ANTHROPIC, + endpoint_type=EndpointType(endpoint_type), start_time=start, raw_bytes=chunks, end_time=end, @@ -374,14 +376,18 @@ def stream_success( def stream_failure( - logger: Logging, request_body: dict[str, object], chunks: list[bytes], error: Exception + logger: Logging, + endpoint_type: str, + request_body: dict[str, object], + chunks: list[bytes], + error: Exception, ) -> Coroutine[object, object, None]: from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType return PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=logger, - endpoint_type=EndpointType.ANTHROPIC, + endpoint_type=EndpointType(endpoint_type), request_body=request_body, raw_bytes=chunks, exception=error, From 19ffb584eb6b9c28118f96fb47642d1fca5442b3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:37:05 -0700 Subject: [PATCH 096/102] refactor(rust): rename litellm-callbacks to litellm-host and HostOpError to InvokeError --- litellm-rust/Cargo.lock | 28 +++++++++---------- litellm-rust/Cargo.toml | 2 +- .../crates/callbacks-legacy/AGENTS.md | 2 +- .../crates/callbacks-legacy/Cargo.toml | 2 +- .../crates/callbacks-legacy/src/adapter.rs | 2 +- .../crates/callbacks-legacy/src/call.rs | 2 +- .../crates/callbacks-legacy/src/callbacks.rs | 2 +- .../tests/deployment_hooks.rs | 2 +- .../crates/callbacks-legacy/tests/payload.rs | 2 +- .../crates/callbacks-legacy/tests/terminal.rs | 2 +- litellm-rust/crates/core/Cargo.toml | 2 +- litellm-rust/crates/core/src/machine/auth.rs | 2 +- litellm-rust/crates/core/src/machine/mod.rs | 2 +- litellm-rust/crates/core/src/messages/mod.rs | 2 +- .../crates/core/src/messages/route.rs | 4 +-- litellm-rust/crates/core/src/ocr/client.rs | 2 +- litellm-rust/crates/core/src/ocr/handler.rs | 2 +- litellm-rust/crates/core/src/ocr/route.rs | 4 +-- .../tests/azure_document_intelligence_ocr.rs | 4 +-- litellm-rust/crates/core/tests/ocr.rs | 6 ++-- .../crates/core/tests/ocr/document.rs | 2 +- litellm-rust/crates/core/tests/ocr/support.rs | 4 +-- litellm-rust/crates/core/tests/reducto_ocr.rs | 4 +-- litellm-rust/crates/host-python/AGENTS.md | 4 +-- litellm-rust/crates/host-python/Cargo.toml | 2 +- .../crates/host-python/src/adapter.rs | 10 +++---- litellm-rust/crates/host-python/src/driver.rs | 26 ++++++++--------- litellm-rust/crates/host-python/src/lib.rs | 4 +-- .../crates/{callbacks => host}/Cargo.toml | 2 +- .../crates/{callbacks => host}/src/event.rs | 0 .../crates/{callbacks => host}/src/host.rs | 0 .../crates/{callbacks => host}/src/lib.rs | 0 .../crates/{callbacks => host}/src/machine.rs | 0 .../crates/{callbacks => host}/src/route.rs | 0 .../crates/{callbacks => host}/src/run.rs | 0 litellm-rust/crates/llms/Cargo.toml | 2 +- .../llms/src/custom_httpx/llm_http_handler.rs | 2 +- .../python-bridge/src/routes/messages/host.rs | 6 ++-- .../python-bridge/src/routes/ocr/host.rs | 6 ++-- 39 files changed, 75 insertions(+), 75 deletions(-) rename litellm-rust/crates/{callbacks => host}/Cargo.toml (91%) rename litellm-rust/crates/{callbacks => host}/src/event.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/host.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/lib.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/machine.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/route.rs (100%) rename litellm-rust/crates/{callbacks => host}/src/run.rs (100%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index bf58a81a6c5..5a8a613204f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2027,22 +2027,12 @@ dependencies = [ "tokio", ] -[[package]] -name = "litellm-callbacks" -version = "0.1.0" -dependencies = [ - "litellm-auth", - "rstest", - "serde_json", - "tokio", -] - [[package]] name = "litellm-callbacks-legacy" version = "0.1.0" dependencies = [ "litellm-auth", - "litellm-callbacks", + "litellm-host", "litellm-host-python", "proptest", "pyo3", @@ -2060,8 +2050,8 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-aws", - "litellm-callbacks", "litellm-core-utils", + "litellm-host", "litellm-llms", "litellm-types", "mime_guess", @@ -2114,12 +2104,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-host" +version = "0.1.0" +dependencies = [ + "litellm-auth", + "rstest", + "serde_json", + "tokio", +] + [[package]] name = "litellm-host-python" version = "0.1.0" dependencies = [ "futures-util", - "litellm-callbacks", + "litellm-host", "pyo3", "pyo3-async-runtimes", "pythonize", @@ -2143,9 +2143,9 @@ dependencies = [ "litellm-auth-aws", "litellm-auth-azure", "litellm-auth-gcp", - "litellm-callbacks", "litellm-core-utils", "litellm-framing", + "litellm-host", "litellm-types", "reqwest 0.12.28", "rstest", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 32f925b8d8b..de6eacc62ee 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -10,7 +10,7 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } -litellm-callbacks = { path = "crates/callbacks" } +litellm-host = { path = "crates/host" } litellm-callbacks-legacy = { path = "crates/callbacks-legacy" } litellm-framing = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } diff --git a/litellm-rust/crates/callbacks-legacy/AGENTS.md b/litellm-rust/crates/callbacks-legacy/AGENTS.md index e184e2fb415..8b2e1c15f6e 100644 --- a/litellm-rust/crates/callbacks-legacy/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy/AGENTS.md @@ -11,7 +11,7 @@ - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup` - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only - - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view + - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-host`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view - Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy/Cargo.toml index efacde051ed..023c13d912b 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true autotests = false [dependencies] -litellm-callbacks.workspace = true +litellm-host.workspace = true litellm-host-python.workspace = true pyo3.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 4e9167100e0..6c013cd1ea5 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -2,7 +2,7 @@ //! raises is answered with the same `Logging` calls, in the same order, as the Python //! `@client` path makes them. -use litellm_callbacks::event::{ +use litellm_host::event::{ FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds, }; use litellm_host_python::{ diff --git a/litellm-rust/crates/callbacks-legacy/src/call.rs b/litellm-rust/crates/callbacks-legacy/src/call.rs index bd1e2525d3d..b37790f60a8 100644 --- a/litellm-rust/crates/callbacks-legacy/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy/src/call.rs @@ -3,7 +3,7 @@ //! lifetime. No other callback host has that obligation, which is why nothing outside //! this crate holds them. -use litellm_callbacks::{machine::Machine, route::Route}; +use litellm_host::{machine::Machine, route::Route}; use litellm_host_python::{RouteHost, lookup, run_call}; use pyo3::{ gc::{PyTraverseError, PyVisit}, diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs index 9fcfe98368e..5f04224e6d7 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs @@ -2,7 +2,7 @@ //! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls //! duplication. All of it expires with the legacy callback contract. -use litellm_callbacks::event::{RequestContext, WireRequest}; +use litellm_host::event::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index f4602739548..52c5e47f83f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,6 +1,6 @@ use std::ffi::CStr; -use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host::event::{FailureOrigin, Timing}; use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 67ad4ab8a2a..5459b36af27 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; use litellm_auth::SecretValue; -use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; use proptest::prelude::*; use pyo3::prelude::*; diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 5688f70f387..f68209233f2 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,6 +1,6 @@ use std::ffi::CStr; -use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host::event::{FailureOrigin, Timing}; use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index db6cfc4b340..3995a235778 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -9,7 +9,7 @@ autotests = false [dependencies] litellm-types.workspace = true litellm-core-utils.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true base64.workspace = true diff --git a/litellm-rust/crates/core/src/machine/auth.rs b/litellm-rust/crates/core/src/machine/auth.rs index 6a3e4daf6ee..cf91458ae43 100644 --- a/litellm-rust/crates/core/src/machine/auth.rs +++ b/litellm-rust/crates/core/src/machine/auth.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; -use litellm_callbacks::route::Route; +use litellm_host::route::Route; use super::{HostChannel, MachineFault}; diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/core/src/machine/mod.rs index d6db488159e..ffcefc663ab 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/core/src/machine/mod.rs @@ -7,7 +7,7 @@ mod auth; use std::{future::Future, pin::Pin}; pub use auth::{HostTokenProvider, TokenRoute}; -use litellm_callbacks::{ +use litellm_host::{ event::{MachineEvent, RequestContext, WireRequest}, host::{Demand, HostOp, HostResult}, machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index e36c6668efe..c07a83c7e43 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -33,7 +33,7 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result Ok(message), MessagesOutput::Streamed => Err(Error::Unsupported( "streamed responses need a streaming host", diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 2fd2f79907a..549f848049f 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -2,12 +2,12 @@ use std::{sync::Mutex, time::Duration}; use bytes::Bytes; use litellm_auth::SecretValue; -use litellm_callbacks::{ +use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; +use litellm_host::{ event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, route::Route, }; -use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 03782d91f24..c05622932b1 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -12,7 +12,7 @@ pub async fn perform( client: &OcrClient, request: LiteLLMOcrRequest, ) -> Result { - litellm_callbacks::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await + litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await } pub async fn ocr(request: LiteLLMOcrRequest) -> Result { diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index aff1eded2cc..bbf9cfa0e02 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,6 +1,6 @@ use futures_util::future::BoxFuture; use litellm_auth::SecretValue; -use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index e6ee45c64a8..d3711c34fa7 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use litellm_auth::ResolvedCredential; -use litellm_callbacks::{ +use litellm_host::{ event::{CallEvent, RequestContext, WireRequest}, route::Route, }; @@ -175,7 +175,7 @@ impl LocalOcrHost { } } -impl litellm_callbacks::host::Host for LocalOcrHost { +impl litellm_host::host::Host for LocalOcrHost { async fn route(&self, op: OcrOp) -> Result { match op { OcrOp::ProjectRequest => self diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 62544931dfc..3cbe6fe3159 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, MachineEvent}; +use litellm_host::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::error::Error; use rstest::rstest; use serde_json::{Value, json}; @@ -466,7 +466,7 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() { mod transformation { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::{CallEvent, MachineEvent}; + use litellm_host::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::transformation::OcrDocument; use serde_json::{Value, json}; diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index a6001951361..41a650945bc 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,6 +1,6 @@ use std::sync::{Arc, Mutex}; -use litellm_callbacks::{ +use litellm_host::{ event::{CallEvent, MachineEvent, WireRequest}, host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, @@ -820,7 +820,7 @@ impl Host for CallerTokenHost { async fn before_send( &self, wire: WireRequest, - _: &litellm_callbacks::event::RequestContext, + _: &litellm_host::event::RequestContext, ) -> Result { let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); let authorization = wire @@ -855,7 +855,7 @@ async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_ trace: Mutex::new(Vec::new()), }; - litellm_callbacks::run::run(ocr_machine(ocr_client()), &host) + litellm_host::run::run(ocr_machine(ocr_client()), &host) .await .unwrap(); server.await.unwrap(); diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs index 5e10ce3ad2a..855548dc6bf 100644 --- a/litellm-rust/crates/core/tests/ocr/document.rs +++ b/litellm-rust/crates/core/tests/ocr/document.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::WireRequest; +use litellm_host::event::WireRequest; use litellm_llms::base_llm::ocr::error::Error; use rstest::rstest; use serde_json::{Value, json}; diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index 44313d5f552..b368a754656 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use futures_util::future::BoxFuture; -use litellm_callbacks::event::WireRequest; +use litellm_host::event::WireRequest; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::{CallHooks, OcrClient}, @@ -45,7 +45,7 @@ pub(crate) async fn perform_ocr(request: LiteLLMOcrRequest) -> Result Result { - litellm_callbacks::run::run(ocr_machine(ocr_client()), &host).await + litellm_host::run::run(ocr_machine(ocr_client()), &host).await } pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 8c037889a17..83e7754122b 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; +use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; use rstest::rstest; use serde_json::{Value, json}; @@ -351,7 +351,7 @@ async fn guardrail_rewrites_document_before_upload() { } mod transformation { - use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::{ base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, reducto::ocr::transformation::*, diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index 903ccb06c77..5aca13eeb18 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -1,9 +1,9 @@ - Target invariants; implementation and runtime validation may lag these rules - Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`RouteHost` traits - - No LiteLLM domain dependencies beyond `litellm-callbacks`: no route types, no `Logging` policy, no public API registration, no cdylib build features + - No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business - `RouteHost::invoke` receives the keyword view the adapter's `begin` returned, not the caller's dict; a route host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) - - A native failure, including one a host op returns as `HostOpError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is + - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index ae0cebada59..e2c83fe1081 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] futures-util.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true pythonize.workspace = true diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index e70e4a8c58f..3a4cb49be4d 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -1,5 +1,5 @@ -use litellm_callbacks::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; -use litellm_callbacks::route::Route; +use litellm_host::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; +use litellm_host::route::Route; use pyo3::exceptions::PyRuntimeError; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -89,12 +89,12 @@ pub trait PythonLifecycle: Send + Sync { /// rejected it, which the route classifies like any other native failure, or Python code /// raised, which reaches the caller as it was raised. #[derive(Debug)] -pub enum HostOpError { +pub enum InvokeError { Native(E), Python(PyErr), } -impl From for HostOpError { +impl From for InvokeError { fn from(error: PyErr) -> Self { Self::Python(error) } @@ -117,7 +117,7 @@ pub trait RouteHost: Send + Sync { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: ::Op, - ) -> Result<::OpResult, HostOpError<::Error>>; + ) -> Result<::OpResult, InvokeError<::Error>>; fn complete( &mut self, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index d6b31b27314..392d36e10f4 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -2,10 +2,10 @@ use std::sync::Arc; use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; -use litellm_callbacks::event::{FailureOrigin, Timing, epoch_seconds}; -use litellm_callbacks::host::{Demand, HostOp, HostResult, HostStep}; -use litellm_callbacks::machine::{HostFailure, Machine, MachineStep}; -use litellm_callbacks::route::Route; +use litellm_host::event::{FailureOrigin, Timing, epoch_seconds}; +use litellm_host::host::{Demand, HostOp, HostResult, HostStep}; +use litellm_host::machine::{HostFailure, Machine, MachineStep}; +use litellm_host::route::Route; use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -13,7 +13,7 @@ use pyo3::types::PyDict; use tokio::sync::Mutex; use crate::adapter::{ - HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, + InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -282,12 +282,12 @@ where let arguments = self.arguments.as_ref().ok_or_else(missing_state)?; match self.route.invoke(py, arguments.bind(py), op) { Ok(result) => Ok(HostResult::Route(result)), - Err(HostOpError::Native(error)) => { + Err(InvokeError::Native(error)) => { return self .resume_core(py, Some(Err(HostFailure::Error(error)))) .map(Next::Continue); } - Err(HostOpError::Python(error)) => Err(error), + Err(InvokeError::Python(error)) => Err(error), } } HostOp::BeforeSend { wire, context } => { @@ -534,8 +534,8 @@ where mod tests { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::{MachineEvent, RequestContext, WireRequest}; - use litellm_callbacks::machine::{Interrupted, Step}; + use litellm_host::event::{MachineEvent, RequestContext, WireRequest}; + use litellm_host::machine::{Interrupted, Step}; use pyo3::exceptions::{PyBaseException, PyValueError}; use pyo3::types::PyDict; @@ -692,12 +692,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri _: Python<'_>, arguments: &Bound<'_, PyDict>, op: &'static str, - ) -> Result> { + ) -> Result> { self.log.push(format!("route:{op}")); match self.op { OpScript::Answer => Ok(format!("{op}:{}", arguments.len())), OpScript::RaisePython => Err(PyValueError::new_err("op failed").into()), - OpScript::RejectNatively => Err(HostOpError::Native(Error("op rejected".into()))), + OpScript::RejectNatively => Err(InvokeError::Native(Error("op rejected".into()))), } } @@ -899,7 +899,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri context: Box::new(context()), }, HostOp::Emit(MachineEvent::ResponseReceived { - raw: litellm_callbacks::event::RawResponse { body: "raw".into() }, + raw: litellm_host::event::RawResponse { body: "raw".into() }, }), ], outcome: Some(Ok("done".into())), @@ -1189,7 +1189,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, _: &Bound<'_, PyDict>, _: &'static str, - ) -> Result> { + ) -> Result> { self.0.push("route"); Err(PyErr::from_value( py.import("asyncio") diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 3738a9b1c3c..583a4eb91b6 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -1,5 +1,5 @@ //! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and -//! asyncio glue, and the driver that runs a native [`Machine`](litellm_callbacks::machine::Machine) +//! asyncio glue, and the driver that runs a native [`Machine`](litellm_host::machine::Machine) //! against a Python route host and a Python lifecycle. Everything here is Python-specific by //! construction; another host language gets its own crate of the same shape. @@ -13,7 +13,7 @@ mod handle; mod marshal; pub use adapter::{ - HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, + InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; pub use argument::lookup; pub use callable::wrap_failure; diff --git a/litellm-rust/crates/callbacks/Cargo.toml b/litellm-rust/crates/host/Cargo.toml similarity index 91% rename from litellm-rust/crates/callbacks/Cargo.toml rename to litellm-rust/crates/host/Cargo.toml index a68ebc26a8d..ebabe5ccbdc 100644 --- a/litellm-rust/crates/callbacks/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "litellm-callbacks" +name = "litellm-host" version = "0.1.0" edition.workspace = true license.workspace = true diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/host/src/event.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/event.rs rename to litellm-rust/crates/host/src/event.rs diff --git a/litellm-rust/crates/callbacks/src/host.rs b/litellm-rust/crates/host/src/host.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/host.rs rename to litellm-rust/crates/host/src/host.rs diff --git a/litellm-rust/crates/callbacks/src/lib.rs b/litellm-rust/crates/host/src/lib.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/lib.rs rename to litellm-rust/crates/host/src/lib.rs diff --git a/litellm-rust/crates/callbacks/src/machine.rs b/litellm-rust/crates/host/src/machine.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/machine.rs rename to litellm-rust/crates/host/src/machine.rs diff --git a/litellm-rust/crates/callbacks/src/route.rs b/litellm-rust/crates/host/src/route.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/route.rs rename to litellm-rust/crates/host/src/route.rs diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/host/src/run.rs similarity index 100% rename from litellm-rust/crates/callbacks/src/run.rs rename to litellm-rust/crates/host/src/run.rs diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 4ca6c7cb2a5..d295e4407ba 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -15,7 +15,7 @@ litellm-auth.workspace = true litellm-auth-aws.workspace = true litellm-auth-azure.workspace = true litellm-auth-gcp.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true litellm-framing.workspace = true base64.workspace = true bytes.workspace = true diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index 7635bdd3d04..fdd568d83fd 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -3,7 +3,7 @@ use std::{sync::OnceLock, time::Duration}; use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; -use litellm_callbacks::event::WireRequest; +use litellm_host::event::WireRequest; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 0c87aea1ca2..b590018bc4e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -3,7 +3,7 @@ use litellm_core::messages::{ Error, route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput}, }; -use litellm_host_python::{HostOpError, RouteHost, from_py, lookup, to_py}; +use litellm_host_python::{InvokeError, RouteHost, from_py, lookup, to_py}; use litellm_llms::custom_httpx::transport::Error as TransportError; use pyo3::{ exceptions::{PyException, PyValueError}, @@ -136,12 +136,12 @@ impl RouteHost for MessagesRouteHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: MessagesOp, - ) -> Result> { + ) -> Result> { match op { MessagesOp::ProjectRequest => self .project(py, arguments) .map(|call| MessagesOpResult::Request(Box::new(call))) - .map_err(|error| HostOpError::Python(self.map_failure(py, error))), + .map_err(|error| InvokeError::Python(self.map_failure(py, error))), } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 19211671fd7..77c8d5d6641 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,6 +1,6 @@ use litellm_auth::ResolvedCredential; use litellm_core::ocr::route::{Ocr, OcrOp, OcrOpResult}; -use litellm_host_python::{HostOpError, RouteHost, missing_state, to_py}; +use litellm_host_python::{InvokeError, RouteHost, missing_state, to_py}; use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -113,9 +113,9 @@ impl RouteHost for OcrRouteHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: OcrOp, - ) -> Result> { + ) -> Result> { self.answer(py, arguments, op) - .map_err(|error| HostOpError::Python(self.map_failure(py, error))) + .map_err(|error| InvokeError::Python(self.map_failure(py, error))) } fn complete(&mut self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult> { From b7686d78b675a020bb15123c571438f647674431 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:40:02 -0700 Subject: [PATCH 097/102] refactor(rust): move RouteMachine out of core into litellm-host next to the Machine trait core/src/machine was route-neutral runtime code sitting among route surfaces, and the workspace had two modules named machine. It now lives in litellm-host beside the contract it implements, so core only holds routes. The OCR conversion from MachineFault moves to litellm-llms because the orphan rule no longer allows it in core --- litellm-rust/crates/core/src/lib.rs | 1 - litellm-rust/crates/core/src/messages/route.rs | 6 ++---- litellm-rust/crates/core/src/ocr/route.rs | 16 ++-------------- litellm-rust/crates/host/Cargo.toml | 2 +- litellm-rust/crates/host/src/lib.rs | 2 +- .../crates/{core => host}/src/machine/auth.rs | 5 ++--- .../host/src/{machine.rs => machine/mod.rs} | 6 ++++++ .../mod.rs => host/src/machine/route_machine.rs} | 10 ++++------ .../crates/llms/src/base_llm/ocr/error.rs | 11 +++++++++++ 9 files changed, 29 insertions(+), 30 deletions(-) rename litellm-rust/crates/{core => host}/src/machine/auth.rs (98%) rename litellm-rust/crates/host/src/{machine.rs => machine/mod.rs} (91%) rename litellm-rust/crates/{core/src/machine/mod.rs => host/src/machine/route_machine.rs} (97%) diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 58aef6cd629..e3e2fb48721 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -2,7 +2,6 @@ pub mod audio_transcription; pub mod chat_completions; pub mod constants; pub mod error; -pub mod machine; pub mod messages; pub mod ocr; pub mod responses; diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 549f848049f..3d607ba67f6 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -6,6 +6,7 @@ use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{ event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, + machine::{HostChannel, MachineFault, RouteMachine}, route::Route, }; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; @@ -18,10 +19,7 @@ use super::{ prepare::prepare_provider_request, types::MessagesRequest, }; -use crate::{ - constants::ANTHROPIC_MESSAGES_PROVIDER, - machine::{HostChannel, MachineFault, RouteMachine}, -}; +use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MessagesOp { diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index d3711c34fa7..bfc8c5ca965 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -3,6 +3,7 @@ use std::sync::{Arc, Mutex}; use litellm_auth::ResolvedCredential; use litellm_host::{ event::{CallEvent, RequestContext, WireRequest}, + machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute}, route::Route, }; use litellm_llms::{ @@ -11,10 +12,7 @@ use litellm_llms::{ }; use super::handler::perform_ocr_request; -use crate::{ - machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute}, - ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest}, -}; +use crate::ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum OcrOp { @@ -56,16 +54,6 @@ impl TokenRoute for Ocr { } } -impl From for Error { - fn from(fault: MachineFault) -> Self { - Self::InvalidRequest(match fault { - MachineFault::Abandoned => "OCR host driver was abandoned".into(), - MachineFault::Protocol(message) => format!("OCR {message}"), - MachineFault::Mismatch => "invalid OCR host operation result".into(), - }) - } -} - pub type OcrHost = HostChannel; pub type OcrMachine = RouteMachine; diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index ebabe5ccbdc..0c7c46192b5 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -8,7 +8,7 @@ repository.workspace = true [dependencies] litellm-auth.workspace = true serde_json.workspace = true +tokio = { workspace = true, features = ["sync"] } [dev-dependencies] rstest.workspace = true -tokio = { workspace = true, features = ["macros"] } diff --git a/litellm-rust/crates/host/src/lib.rs b/litellm-rust/crates/host/src/lib.rs index 41b0983f0ce..65479c2380f 100644 --- a/litellm-rust/crates/host/src/lib.rs +++ b/litellm-rust/crates/host/src/lib.rs @@ -1,7 +1,7 @@ //! The contract between a native call and the host runtime that drives it. //! //! A host is whatever sits on the far side of the language boundary: CPython today, -//! another runtime later. Core implements [`machine::Machine`] per route and never learns +//! another runtime later. Core runs each route on a [`machine::RouteMachine`] and never learns //! which host is on the other end. The machine yields [`host::HostOp`]s; a driver answers //! them, observes [`event::CallEvent`]s and may rewrite the wire request before it is sent. diff --git a/litellm-rust/crates/core/src/machine/auth.rs b/litellm-rust/crates/host/src/machine/auth.rs similarity index 98% rename from litellm-rust/crates/core/src/machine/auth.rs rename to litellm-rust/crates/host/src/machine/auth.rs index cf91458ae43..ba7e242e766 100644 --- a/litellm-rust/crates/core/src/machine/auth.rs +++ b/litellm-rust/crates/host/src/machine/auth.rs @@ -1,9 +1,8 @@ use std::sync::Arc; -use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; -use litellm_host::route::Route; - use super::{HostChannel, MachineFault}; +use crate::route::Route; +use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; /// A route whose host can mint credentials on the call's behalf. pub trait TokenRoute: Route { diff --git a/litellm-rust/crates/host/src/machine.rs b/litellm-rust/crates/host/src/machine/mod.rs similarity index 91% rename from litellm-rust/crates/host/src/machine.rs rename to litellm-rust/crates/host/src/machine/mod.rs index 2942913f095..2c26db61582 100644 --- a/litellm-rust/crates/host/src/machine.rs +++ b/litellm-rust/crates/host/src/machine/mod.rs @@ -1,6 +1,12 @@ +mod auth; +mod route_machine; + use std::future::Future; use std::pin::Pin; +pub use auth::{HostTokenProvider, TokenRoute}; +pub use route_machine::{ExecuteFuture, HostChannel, MachineFault, RouteMachine}; + use crate::host::{HostOp, HostResult}; use crate::route::Route; diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/host/src/machine/route_machine.rs similarity index 97% rename from litellm-rust/crates/core/src/machine/mod.rs rename to litellm-rust/crates/host/src/machine/route_machine.rs index ffcefc663ab..38a0b8bc16a 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/host/src/machine/route_machine.rs @@ -2,18 +2,16 @@ //! place, and turns the host operations that future requests into [`Machine`] steps. No //! task is spawned; dropping the machine drops the in-flight call. -mod auth; - use std::{future::Future, pin::Pin}; -pub use auth::{HostTokenProvider, TokenRoute}; -use litellm_host::{ +use tokio::sync::{mpsc, oneshot}; + +use super::{HostFailure, Interrupted, Machine, MachineStep, Step}; +use crate::{ event::{MachineEvent, RequestContext, WireRequest}, host::{Demand, HostOp, HostResult}, - machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, route::Route, }; -use tokio::sync::{mpsc, oneshot}; /// The machine's own failures, distinct from anything the provider call reports. #[derive(Clone, Copy, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index c3f481d7d44..3061a9fe2b2 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -102,6 +102,17 @@ pub enum Error { Headers(#[from] crate::custom_httpx::http_handler::HeaderError), } +impl From for Error { + fn from(fault: litellm_host::machine::MachineFault) -> Self { + use litellm_host::machine::MachineFault; + Self::InvalidRequest(match fault { + MachineFault::Abandoned => "OCR host driver was abandoned".into(), + MachineFault::Protocol(message) => format!("OCR {message}"), + MachineFault::Mismatch => "invalid OCR host operation result".into(), + }) + } +} + impl From for Error { fn from(error: litellm_core_utils::call_arguments::ArgumentError) -> Self { Self::RequestField { From 41873e2bc796bce93b74ec9431ae01b403ed46a3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 22:41:17 +0000 Subject: [PATCH 098/102] fix(rust): box messages response output Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/core/src/messages/mod.rs | 2 +- litellm-rust/crates/core/src/messages/route.rs | 5 +++-- .../crates/python-bridge/src/routes/messages/host.rs | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index c07a83c7e43..289f79109dd 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -34,7 +34,7 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result Ok(message), + MessagesOutput::Message(message) => Ok(*message), MessagesOutput::Streamed => Err(Error::Unsupported( "streamed responses need a streaming host", )), diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 3d607ba67f6..838b56fcb4b 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -48,7 +48,7 @@ impl MessagesCall { } pub enum MessagesOutput { - Message(AnthropicMessagesResponse), + Message(Box), /// Every chunk already reached the host through `Deliver`. Streamed, } @@ -174,7 +174,8 @@ async fn execute(host: MessagesHost) -> Result { raw: RawResponse { body: text.clone() }, }) .await?; - decode_response(request.config, &request.model, &text).map(MessagesOutput::Message) + decode_response(request.config, &request.model, &text) + .map(|message| MessagesOutput::Message(Box::new(message))) } /// Hands each upstream chunk to the caller as it arrives. A caller that stops reading diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index b590018bc4e..c1b3f59df58 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -150,7 +150,7 @@ impl RouteHost for MessagesRouteHost { MessagesOutput::Message(message) => py .import("litellm.rust_bridge.messages.route_host")? .getattr("response")? - .call1((to_py(py, &message)?,)) + .call1((to_py(py, message.as_ref())?,)) .map(Bound::unbind), MessagesOutput::Streamed => Ok(py.None()), } From ba6b22cf56ead7aba23d530b342da6d342230581 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:57:50 -0700 Subject: [PATCH 099/102] test(rust): isolate callback registries per hypothesis example Replace the module-level LATEST_EDITS list with per-example callback registry isolation, and import litellm names with from-imports in the legacy callback shim so the module uses one import style. Co-Authored-By: Claude Opus 5 --- litellm/rust_bridge/legacy_callbacks.py | 18 +++--- tests/test_litellm_rust/conftest.py | 63 +++---------------- tests/test_litellm_rust/ocr/test_callbacks.py | 20 +++--- tests/test_litellm_rust/support/isolation.py | 58 +++++++++++++++++ 4 files changed, 88 insertions(+), 71 deletions(-) create mode 100644 tests/test_litellm_rust/support/isolation.py diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index bac40442ce5..30aa1d97bfc 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -65,13 +65,17 @@ def setup( def check_limits(kwargs: Mapping[str, object]) -> None: - import litellm + from litellm import ( + BudgetExceededError, + _current_cost, # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor + max_budget, + num_retries_per_request, + ) from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit - current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor - if litellm.max_budget and current_cost > litellm.max_budget: - raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget) - if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + if max_budget and _current_cost > max_budget: + raise BudgetExceededError(current_cost=_current_cost, max_budget=max_budget) + if max_retries_per_request_hit(kwargs, num_retries_per_request): raise RuntimeError("Max retries per request hit!") @@ -281,9 +285,9 @@ def is_internal_call() -> bool: def credential_list() -> list[CredentialItem]: - import litellm + from litellm import credential_list as credentials - return litellm.credential_list + return credentials def warn_unknown_credential(name: str, loaded: int) -> None: diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 4387ea2e2fd..1b6fcfa00db 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,10 +1,9 @@ import asyncio import os -from collections.abc import AsyncIterator, Generator, Iterator +from collections.abc import AsyncIterator, Generator from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack, contextmanager -from types import ModuleType -from typing import Final, cast +from contextlib import ExitStack +from typing import Final import pytest import pytest_asyncio @@ -18,62 +17,20 @@ from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivate _parse_env_bool, ) from tests.test_litellm_rust.support.callback_recorder import drain_logging +from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service -CALLBACK_ATTRIBUTES: Final = ( - "callbacks", - "input_callback", - "success_callback", - "failure_callback", - "_async_input_callback", - "_async_success_callback", - "_async_failure_callback", -) - - -def _list_attribute(container: ModuleType, attribute: str) -> list[object]: - value: Final = getattr(container, attribute) - if not isinstance(value, list): - raise AssertionError(f"{container.__name__}.{attribute} is not a list") - return cast(list[object], value) - - -@contextmanager -def _isolated_list(container: ModuleType, attribute: str) -> Iterator[None]: - source: Final = _list_attribute(container, attribute) - original: Final = list(source) - source.clear() # mutable-ok: test isolation mutates global registries by design - try: - yield - finally: - source.clear() - source.extend(original) - setattr(container, attribute, source) - - -@contextmanager -def _rebound(container: object, attribute: str, value: object) -> Iterator[None]: - original: Final[object] = getattr(container, attribute) - setattr(container, attribute, value) - try: - yield - finally: - setattr(container, attribute, original) - @pytest_asyncio.fixture(autouse=True, loop_scope="function") async def isolate_ocr_test_state() -> AsyncIterator[None]: with ExitStack() as stack: - for attribute in CALLBACK_ATTRIBUTES: - stack.enter_context(_isolated_list(litellm, attribute)) - stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor - stack.enter_context(_rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry - stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache - stack.enter_context(_rebound(_CONFIGURATION, "override", None)) + stack.enter_context(isolated_callback_registries()) + stack.enter_context(rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache + stack.enter_context(rebound(_CONFIGURATION, "override", None)) executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") - stack.enter_context(_rebound(litellm_logging, "executor", executor)) - stack.enter_context(_rebound(utils, "executor", executor)) - stack.enter_context(_rebound(thread_pool_executor, "executor", executor)) + stack.enter_context(rebound(litellm_logging, "executor", executor)) + stack.enter_context(rebound(utils, "executor", executor)) + stack.enter_context(rebound(thread_pool_executor, "executor", executor)) try: yield finally: diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 45b99d19d90..ac4a1a11a80 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -3,6 +3,8 @@ import copy import gc import queue import threading +from collections.abc import Mapping +from types import MappingProxyType from typing import Final import pytest @@ -13,6 +15,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.isolation import isolated_callback_registries from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, @@ -307,18 +310,13 @@ JSON_VALUES: Final = st.recursive( ) -LATEST_EDITS: Final[list[dict[str, object]]] = [] - - -class ApplyLatestEdits(CustomLogger): - """Registrations can outlive one hypothesis example, so every instance applies the current example's edits.""" - - def __init__(self, latest: list[dict[str, object]]) -> None: +class ApplyEdits(CustomLogger): + def __init__(self, edits: Mapping[str, object]) -> None: super().__init__() - self.latest = latest + self.edits: Final = edits def log_pre_api_call(self, model, messages, kwargs): - request_body(kwargs).update(copy.deepcopy(self.latest[-1])) + request_body(kwargs).update(copy.deepcopy(dict(self.edits))) @settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) @@ -327,9 +325,9 @@ def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_lef ocr_server: RecordingServer, edits: dict[str, object] ) -> None: ocr_server.expected_requests = None - LATEST_EDITS.append(edits) - call_native_ocr_with_callbacks(ocr_server, [ApplyLatestEdits(LATEST_EDITS)]) + with isolated_callback_registries(): + call_native_ocr_with_callbacks(ocr_server, [ApplyEdits(MappingProxyType(edits))]) assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits} diff --git a/tests/test_litellm_rust/support/isolation.py b/tests/test_litellm_rust/support/isolation.py new file mode 100644 index 00000000000..f98ce4843a8 --- /dev/null +++ b/tests/test_litellm_rust/support/isolation.py @@ -0,0 +1,58 @@ +from collections.abc import Generator +from contextlib import ExitStack, contextmanager +from types import ModuleType +from typing import Final, cast + +import litellm +from litellm import utils +from litellm.litellm_core_utils import litellm_logging + +CALLBACK_ATTRIBUTES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +def _list_attribute(container: ModuleType, attribute: str) -> list[object]: + value: Final = getattr(container, attribute) + if not isinstance(value, list): + raise AssertionError(f"{container.__name__}.{attribute} is not a list") + return cast(list[object], value) + + +@contextmanager +def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]: + source: Final = _list_attribute(container, attribute) + original: Final = list(source) + source.clear() # mutable-ok: test isolation mutates global registries by design + try: + yield + finally: + source.clear() + source.extend(original) + setattr(container, attribute, source) + + +@contextmanager +def rebound(container: object, attribute: str, value: object) -> Generator[None]: + original: Final[object] = getattr(container, attribute) + setattr(container, attribute, value) + try: + yield + finally: + setattr(container, attribute, original) + + +@contextmanager +def isolated_callback_registries() -> Generator[None]: + with ExitStack() as stack: + for attribute in CALLBACK_ATTRIBUTES: + stack.enter_context(_isolated_list(litellm, attribute)) + stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor + stack.enter_context(rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry + yield From 139445179a8a2fd1804cbeba1154e77b4e044fb6 Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 18 Sep 2026 22:59:32 +0000 Subject: [PATCH 100/102] ci: remove the dead Agent Shin triage workflows and scripts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/_agent_shin_actions.py | 50 - .github/scripts/agent_shin_shared.py | 211 -- .github/scripts/close_low_quality_prs.py | 573 ----- .github/scripts/triage-requirements.txt | 282 --- .github/scripts/triage_with_llm.py | 1797 -------------- .github/workflows/close_low_quality_prs.yml | 92 - .../create_daily_oss_agent_shin_branch.yml | 28 - .github/workflows/triage_reconsider.yml | 172 -- .../test_github_close_low_quality_prs.py | 856 ------- tests/test_litellm/test_github_review_gate.py | 524 ---- .../test_github_triage_with_llm.py | 2134 ----------------- .../test_github_triage_workflows.py | 264 -- 12 files changed, 6983 deletions(-) delete mode 100644 .github/scripts/_agent_shin_actions.py delete mode 100644 .github/scripts/agent_shin_shared.py delete mode 100644 .github/scripts/close_low_quality_prs.py delete mode 100644 .github/scripts/triage-requirements.txt delete mode 100644 .github/scripts/triage_with_llm.py delete mode 100644 .github/workflows/close_low_quality_prs.yml delete mode 100644 .github/workflows/create_daily_oss_agent_shin_branch.yml delete mode 100644 .github/workflows/triage_reconsider.yml delete mode 100644 tests/test_litellm/test_github_close_low_quality_prs.py delete mode 100644 tests/test_litellm/test_github_review_gate.py delete mode 100644 tests/test_litellm/test_github_triage_with_llm.py delete mode 100644 tests/test_litellm/test_github_triage_workflows.py diff --git a/.github/scripts/_agent_shin_actions.py b/.github/scripts/_agent_shin_actions.py deleted file mode 100644 index b3d1ff055b3..00000000000 --- a/.github/scripts/_agent_shin_actions.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Dry-run wrapper(s) around Agent Shin GitHub mutations. - -The rollout scripts currently need only one mutation wrapped, so this module -exposes a single ``maybe_post_comment`` helper. It takes a ``dry_run: bool`` -keyword argument and the body is intentionally trivial: - - if dry_run: - print(...) # log what we would do, return - return - real_mutation(...) # otherwise, actually do it - -That shape means a dry-run preview differs from the real run in exactly one -line per side effect: the call site. So when you `python3 script.py` locally -without ``--close``, you can be confident the actions printed are the ones the -GitHub Action would have performed (modulo ordering on retry/error paths, -which are deliberately simple). Any further mutation a rollout script needs -should get the same ``maybe_*`` treatment instead of calling the raw -``triage_with_llm`` mutation directly. - -Importing from this module pulls in the real mutation from ``triage_with_llm`` -— call sites in the rollout scripts should NEVER import ``post_comment`` -directly; that would skip the dry-run gate and is the bug class this module -exists to prevent. -""" - -from __future__ import annotations - -import sys -import textwrap - -# Import the module itself rather than the bare names so monkeypatching -# `triage_with_llm.post_comment` (or any of the other mutations) in tests is -# reflected here — `from triage_with_llm import post_comment` would bind the -# original function to a local name and bypass the patch, defeating the whole -# point of these wrappers. -import triage_with_llm - - -def _log(line: str) -> None: - """Print a single dry-run line to stdout (one log statement per side effect).""" - print(line, file=sys.stdout, flush=True) - - -def maybe_post_comment(repo: str, number: int, body: str, *, dry_run: bool) -> None: - """Post a comment on ``repo#number`` — or, in dry-run, log what we would post.""" - if dry_run: - _log(f"[DRY RUN] comment {repo}#{number}:") - _log(textwrap.indent(body, " ")) - return - triage_with_llm.post_comment(repo, number, body) diff --git a/.github/scripts/agent_shin_shared.py b/.github/scripts/agent_shin_shared.py deleted file mode 100644 index 8f3dc3c2322..00000000000 --- a/.github/scripts/agent_shin_shared.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Constants and helpers shared by Agent Shin's triage scripts. - -Both `triage_with_llm.py` (the LLM-judge entrypoint) and -`close_low_quality_prs.py` (the daily Greptile-score sweep) need to -agree on the same notions of: - - * What counts as a Greptile-authored review comment - (``GREPTILE_BOT_LOGINS``) and how to extract a confidence score from - its body (``SCORE_PATTERN`` / :func:`extract_greptile_score`). - * How long the 2-hour grace window is (``GRACE_PERIOD_SECONDS``) and - the HTML marker stamped into a grace-warning comment so the *other* - script can see "Agent Shin already warned" and behave accordingly - (``GRACE_COMMENT_MARKER``). - * Who Agent Shin is on GitHub (``AGENT_SHIN_DEFAULT_BOT_LOGIN``). - * How GitHub-style ISO-8601 timestamps round-trip into timezone-aware - :class:`datetime.datetime` (:func:`parse_iso8601`). - -Keeping these in one module means a future change (new Greptile output -format, a longer grace window, a new allowlisted account) is a single edit -instead of two — the original split version had to call out in comments -that the two copies "must stay in sync" precisely because nothing -enforced it. -""" - -from __future__ import annotations - -import datetime as dt -import json -import os -import re -import subprocess -from typing import Iterable - -GREPTILE_BOT_LOGINS = frozenset({"greptile-apps", "greptile-apps[bot]"}) - -SCORE_PATTERN = re.compile( - r"confidence\s*score\s*[:\-]?\s*(\d+)\s*/\s*5", - re.IGNORECASE, -) - -GRACE_COMMENT_MARKER = "" - -# Hidden HTML marker stamped on every Agent Shin auto-close comment (the LLM -# judge's grace/review-gate close and the daily Greptile sweep's close). -# `was_closed_by_agent_shin` requires this marker — not just the closing actor — -# before `@agent-shin reconsider` may reopen, because the `github-actions[bot]` -# identity is shared with every other workflow in the repo and is not unique to -# Agent Shin. Both close paths must stamp it or the reconsider path silently -# rejects the contributor. -AGENT_SHIN_CLOSE_MARKER = "" - -# 2 hours between the grace warning and the auto-close. Short enough to -# dogfood the "fix it before it closes" loop in one sitting; bump back up -# (e.g. 86400 for a day) for the public rollout. -GRACE_PERIOD_SECONDS = 7200 - -AGENT_SHIN_DEFAULT_BOT_LOGIN = "github-actions[bot]" - - -def _logins(*names: str) -> frozenset[str]: - """Build a login set normalized for case-insensitive membership checks. - - Callers compare via ``login.lower() in ``, so the stored values - must be lowercase. Normalizing here lets the literals keep each - account's canonical GitHub casing (e.g. ``SwiftWinds``) for - readability without breaking the lookup. - """ - return frozenset(name.lower() for name in names) - - -# Dogfood rollout gate. While this set is non-empty, Agent Shin acts ONLY on -# PRs/issues authored by these logins and skips everyone else. For an -# allowlisted author the usual internal/external classification is bypassed, so -# an internal account (e.g. a maintainer's own work login) still gets triaged -# while the bot is being tested on a small set of accounts. Empty the set to -# lift the restriction and restore full triage for the public rollout. Logins -# are compared case-insensitively. -ALLOWLIST_LOGINS = _logins("mateo-berri", "SwiftWinds") - -# `gh {pr,issue} list` has no "fetch everything" flag — `--limit` is the only -# control and it defaults to 30. Pass a ceiling far above any realistic open -# backlog (low thousands today) so gh paginates the API until the queue is -# exhausted rather than silently truncating. The bulk sweeps MUST see the whole -# backlog: gh lists newest-first, so a low cap drops the *oldest* PRs/issues — -# exactly the stale ones a low-quality sweep is meant to catch. -GH_LIST_ALL_LIMIT = 100_000 - - -def extract_greptile_score(comments: Iterable[dict]) -> tuple[int, dict] | None: - """Return (score, comment) for the most recent Greptile-authored comment - that contains a "Confidence Score: X/5". Returns None if no such comment. - - "Most recent" is determined by the comment's `updated_at` (falling back to - `created_at`), so re-reviews override earlier passes. - """ - candidates: list[tuple[str, int, dict]] = [] - for comment in comments: - user = (comment.get("user") or {}).get("login", "") - if user not in GREPTILE_BOT_LOGINS: - continue - body = comment.get("body") or "" - match = SCORE_PATTERN.search(body) - if not match: - continue - score = int(match.group(1)) - timestamp = comment.get("updated_at") or comment.get("created_at") or "" - candidates.append((timestamp, score, comment)) - - if not candidates: - return None - - candidates.sort(key=lambda triple: triple[0]) - _, score, comment = candidates[-1] - return score, comment - - -def parse_iso8601(value: str) -> dt.datetime: - """Parse a GitHub ISO-8601 timestamp into a timezone-aware datetime.""" - return dt.datetime.fromisoformat(value.replace("Z", "+00:00")) - - -def gh(*args: str) -> str: - """Run a `gh` CLI command and return stdout. Raises on non-zero exit. - - Shared by both Agent Shin entrypoints so a future change here - (timeout handling, logging, retry on transient failures) only needs - to be made once. - """ - result = subprocess.run( - ["gh", *args], - capture_output=True, - text=True, - check=True, - ) - return result.stdout - - -def list_open_items(kind: str, *, repo: str | None, fields: str) -> list[dict]: - """Return EVERY open PR (``kind="pr"``) or issue (``kind="issue"``) in ``repo``. - - Wraps ``gh {pr,issue} list`` with ``--limit GH_LIST_ALL_LIMIT`` so the full - backlog is fetched instead of the default 30 (or any other arbitrary cap). - Both bulk sweeps — the daily Greptile closer and the one-shot rollout - heads-up — rely on this seeing the whole queue, including the oldest items. - - ``fields`` is the comma-separated ``--json`` field list the caller needs - (e.g. ``"number"`` for the rollout, the full set for the closer). - """ - if kind not in ("pr", "issue"): - raise ValueError(f"kind must be 'pr' or 'issue', got {kind!r}") - repo_args = ["--repo", repo] if repo else [] - raw = gh( - kind, - "list", - "--state", - "open", - "--limit", - str(GH_LIST_ALL_LIMIT), - "--json", - fields, - *repo_args, - ) - return json.loads(raw) - - -def seconds_since_latest_marker_comment( - comments: Iterable[dict], - *, - marker: str, - bot_login: str | None = None, - now: dt.datetime | None = None, -) -> float | None: - """Return seconds since the bot's most recent comment containing ``marker``. - - Filters comments by author so a contributor who quotes the HTML - marker (e.g. via GitHub's "Quote reply" feature, which preserves - HTML comments in the raw markdown of the quoted text) is not - mistaken for a bot warning — that would silently reset cooldown - timers and suppress legitimate notifications. - - ``bot_login`` defaults to the `AGENT_SHIN_BOT_LOGIN` env override or - ``AGENT_SHIN_DEFAULT_BOT_LOGIN`` so callers normally don't need to - pass it. ``now`` is injectable for tests / callers (like the daily - sweep) that want every age calculation pinned to one snapshot. - """ - expected_login = ( - bot_login - or os.environ.get("AGENT_SHIN_BOT_LOGIN") - or AGENT_SHIN_DEFAULT_BOT_LOGIN - ).lower() - latest: dt.datetime | None = None - for comment in comments: - author = ((comment.get("user") or {}).get("login") or "").lower() - if author != expected_login: - continue - body = comment.get("body") or "" - if marker not in body: - continue - created = comment.get("created_at") - if not created: - continue - try: - ts = parse_iso8601(created) - except ValueError: - continue - if latest is None or ts > latest: - latest = ts - if latest is None: - return None - reference = now if now is not None else dt.datetime.now(dt.timezone.utc) - return (reference - latest).total_seconds() diff --git a/.github/scripts/close_low_quality_prs.py b/.github/scripts/close_low_quality_prs.py deleted file mode 100644 index 7b9bbb579e3..00000000000 --- a/.github/scripts/close_low_quality_prs.py +++ /dev/null @@ -1,573 +0,0 @@ -#!/usr/bin/env python3 -""" -Auto-close low-quality pull requests. - -Closes open PRs (including drafts, regardless of age) that satisfy ALL of: - 1. Have a Greptile (`greptile-apps`) review comment whose latest - "Confidence Score: X/5" is below the configured threshold (default: 4). - 2. Are authored by an external OSS contributor (internal BerriAI - contributors are exempt). - 3. Do not carry an opt-out label (default: "do not close"). - -`--min-age-days` is retained as an opt-in safety net for one-off backfill -runs (default: 0). The team's intent is that the count of open PRs equals -the count of PRs internal collaborators need to action on, so neither age -nor draft status acts as a free pass. - -For each match, the script posts an explanatory comment and closes the PR. -Because OSS contributors *cannot* reopen a PR closed by the bot/maintainer -(GitHub limitation), the close-comment instructs them to push their fixes -and **open a fresh PR**, or to comment `@agent-shin reconsider` on the -closed PR to have the LLM judge re-evaluate (and reopen on pass). - -Requires the `gh` CLI to be authenticated. - -Usage examples: - # Dry run (default) - prints what would be closed - python3 close_low_quality_prs.py - - # Actually close matching PRs - python3 close_low_quality_prs.py --close - - # Restrict to PRs at least N days old (one-off backfill safety net) - python3 close_low_quality_prs.py --min-age-days 7 --min-score 4 --close -""" - -from __future__ import annotations - -import argparse -import datetime as dt -import json -import os -import subprocess -import sys -from typing import Iterable - -# Add this script's directory to `sys.path` so the sibling -# `agent_shin_shared` module is importable when the script is invoked -# directly (e.g. `python3 .github/scripts/close_low_quality_prs.py ...`). -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from agent_shin_shared import ( # noqa: E402 -- sys.path adjusted above - AGENT_SHIN_CLOSE_MARKER, - ALLOWLIST_LOGINS, - GRACE_COMMENT_MARKER, - GRACE_PERIOD_SECONDS, - GREPTILE_BOT_LOGINS, - SCORE_PATTERN, - extract_greptile_score, - gh, - list_open_items, - parse_iso8601, - seconds_since_latest_marker_comment, -) - -# `GREPTILE_BOT_LOGINS` and `SCORE_PATTERN` (Greptile's GitHub App login -# variants and the "Confidence Score: X/5" regex) are imported from -# `agent_shin_shared` so the LLM judge in `triage_with_llm.py` and this -# daily Greptile sweep read the score through the same set of logins -# and the same regex. - -# `author_association` values for internal BerriAI contributors who should be -# exempt from auto-triage. -INTERNAL_AUTHOR_ASSOCIATIONS = frozenset({"OWNER", "MEMBER", "COLLABORATOR"}) - -# Default labels that exempt a PR from auto-close. Defined at module scope (not -# as a mutable argparse default) so that `--optout-label foo` REPLACES the -# defaults instead of appending to them — the argparse `action="append"` + -# `default=[...]` combination silently mutates the shared default list. -DEFAULT_OPTOUT_LABELS = ("do not close", "keep open", "wip") - -# `GRACE_COMMENT_MARKER` (HTML marker appended to grace-period warning -# comments — used by either script to recognize that a warning was -# already posted) and `GRACE_PERIOD_SECONDS` (length of the grace -# period between the warning and the actual auto-close, 2 hours) are -# imported from `agent_shin_shared` so the Agent Shin LLM judge and -# this daily Greptile sweep agree on the same marker and duration. - - -def fetch_open_prs(repo: str | None) -> list[dict]: - """Fetch all open PRs (number, createdAt, isDraft, labels, author). - - Includes drafts: `gh pr list --state open` returns both ready-for-review - and draft PRs by default. This is the desired behavior — drafts are not - a free pass; the internal-collaborator open-PR queue should reflect every - PR that needs human attention regardless of draft status. - """ - fields = "number,title,createdAt,isDraft,labels,author,url" - return list_open_items("pr", repo=repo, fields=fields) - - -def fetch_pr_author_association(pr_number: int, repo: str | None) -> str: - """Return the GitHub `author_association` for a PR, uppercase. - - Values: OWNER, MEMBER, COLLABORATOR, CONTRIBUTOR, FIRST_TIME_CONTRIBUTOR, - FIRST_TIMER, MANNEQUIN, NONE. Returns "" on lookup failure. - """ - endpoint = ( - f"repos/{repo}/pulls/{pr_number}" - if repo - else f"repos/{{owner}}/{{repo}}/pulls/{pr_number}" - ) - try: - data = json.loads(gh("api", endpoint)) - except subprocess.CalledProcessError: - return "" - return (data.get("author_association") or "").upper() - - -def is_external_pr_author(pr: dict, repo: str | None) -> bool: - """Return True if the PR author is an external OSS contributor. - - Internal = `OWNER` / `MEMBER` / `COLLABORATOR` association, or a bot login. - """ - login = ((pr.get("author") or {}).get("login") or "").lower() - if login.endswith("[bot]") or login in {"dependabot", "github-actions"}: - return False - association = fetch_pr_author_association(pr["number"], repo) - # Fail-safe: if the API lookup failed (empty string), treat the author as - # internal so we don't auto-close their PR. Auto-close is destructive, so - # an unknown association should never make a PR eligible for closing. - if not association or association in INTERNAL_AUTHOR_ASSOCIATIONS: - return False - return True - - -def fetch_pr_comments(pr_number: int, repo: str | None) -> list[dict]: - """Fetch issue-level comments on a PR (where Greptile posts its summary).""" - endpoint = ( - f"repos/{repo}/issues/{pr_number}/comments?per_page=100" - if repo - else f"repos/{{owner}}/{{repo}}/issues/{pr_number}/comments?per_page=100" - ) - raw = gh("api", "--paginate", endpoint) - comments: list[dict] = [] - for line in raw.strip().splitlines(): - line = line.strip() - if not line: - continue - try: - parsed = json.loads(line) - except json.JSONDecodeError: - # A malformed line should not blow up the whole sweep. Skip and - # carry on so the remaining PRs in this run still get evaluated. - continue - if isinstance(parsed, list): - comments.extend(parsed) - else: - comments.append(parsed) - return comments - - -def has_optout_label(pr: dict, optout_labels: set[str]) -> bool: - labels = {label.get("name", "").lower() for label in pr.get("labels", [])} - return bool(labels & {lbl.lower() for lbl in optout_labels}) - - -def seconds_since_last_grace_warning( - comments: Iterable[dict], - *, - bot_login: str | None = None, - now: dt.datetime | None = None, -) -> float | None: - """Return seconds since the bot's most recent grace-period warning, or - None if no such warning has ever been posted on this PR. - - Thin wrapper over - `agent_shin_shared.seconds_since_latest_marker_comment` — the - centralized helper handles the bot-author filter, marker match, - timestamp parsing, and `now` injection. Keeping this wrapper - preserves the closer's "already-fetched comments + injectable now" - interface so callers (and tests) don't need to change. - """ - return seconds_since_latest_marker_comment( - comments, - marker=GRACE_COMMENT_MARKER, - bot_login=bot_login, - now=now, - ) - - -def format_grace_warning_comment(score: int, threshold: int) -> str: - """Comment posted on the FIRST low-Greptile-score detection — gives - the contributor a 2-hour grace window before the auto-close fires on - the next daily cron run. - - Mirrors `format_grace_warning_pr_comment` in - `triage_with_llm.py` in spirit (2-hour grace + escape hatches), but - framed around Greptile's confidence score instead of the LLM judge's - rubric since the close trigger here is the Greptile signal. - """ - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository.\n" - "\n" - "Heads up: Greptile's most recent review scored this PR " - f"**{score}/5**, below our merge bar of **{threshold}/5**.\n" - "\n" - "If the score isn't lifted in the next **2 hours**, I'll auto-close this PR. That's " - "**not** us saying the change isn't worthwhile. We want the open-PR list to mirror " - "what a maintainer can act on *right now*, so contributors like you don't get lost in " - "a backlog. Take your time; everything below still works after the close.\n" - "\n" - "**During the grace period:** push fixes that address Greptile's feedback, then comment " - "`@greptileai` to request a fresh review. If " - f"the new score is **{threshold}/5 or higher**, the PR stays open and no further " - "action is needed on your side.\n" - "\n" - "**If the PR does get auto-closed in 2 hours, you still have an easy recovery path:**\n" - "\n" - "- Comment `@greptileai` to request a fresh review. **This still works even after " - f"the PR is closed**, and a score of {threshold}/5 or higher is one of the signals " - "that lifts the PR back into the review queue. A low Greptile score isn't a blocker.\n" - "- Comment `@agent-shin reconsider` after pushing fixes; I'll re-run the rubric and " - "reopen the PR if both gates (description rubric + Greptile score) now pass.\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -def post_grace_warning( - pr: dict, - score: int, - threshold: int, - repo: str | None, - dry_run: bool, -) -> None: - """Post the 2-hour grace-period warning comment on `pr`. - - The warning carries `GRACE_COMMENT_MARKER` so subsequent runs can - detect that the contributor has already been told about the - pending close. Does NOT close the PR — the close happens on the - next eligible run after `GRACE_PERIOD_SECONDS` elapses (handled - by `close_pr`). - """ - pr_number = pr["number"] - repo_args = ["--repo", repo] if repo else [] - - if dry_run: - print( - f" [DRY RUN] Would post grace warning to PR #{pr_number} " - f"(greptile={score}/5): {pr['title']}" - ) - return - - comment_body = format_grace_warning_comment(score, threshold) - gh("pr", "comment", str(pr_number), "--body", comment_body, *repo_args) - print(f" Posted grace warning on PR #{pr_number} (greptile={score}/5)") - - -def format_close_comment(score: int, threshold: int) -> str: - """Comment posted when a low-Greptile-score PR is auto-closed. - - Carries `AGENT_SHIN_CLOSE_MARKER` so the `@agent-shin reconsider` path - (guarded by `was_closed_by_agent_shin`) recognizes this as an Agent Shin - close and is allowed to reopen the PR once it passes again; without the - marker that recovery path the comment advertises silently rejects the - contributor. - """ - score_sentence = ( - f"Greptile's most recent review scored this PR **{score}/5**, below " - f"our merge bar of **{threshold}/5**, and the 2-hour grace period since " - "the warning has elapsed.\n\n" - ) - return ( - f"Closing as part of automated PR triage.\n\n" - f"{score_sentence}" - "We close low-confidence PRs aggressively to keep the review queue " - "manageable for maintainers and contributors alike. **This is not a " - "rejection of the idea.** To bring this back:\n\n" - "1. Push the fixes that address Greptile's feedback (continue using " - "your existing branch is fine).\n" - "2. **Open a new PR** with the updated branch. Greptile will review " - "it again, and if it scores " - f"**{threshold}/5 or higher** a maintainer will take another look.\n\n" - "_Why open a new PR instead of reopening this one?_ GitHub does not " - "let external contributors reopen a PR that was closed by a bot or " - "maintainer, so a fresh PR is the most reliable path forward. If you " - "would prefer this exact PR re-evaluated, comment " - "`@agent-shin reconsider` once you've pushed the fixes; Agent Shin " - "will re-run triage and reopen this PR if it now meets the bar. " - "You can also comment `@greptileai` to request a fresh Greptile " - "review; that works **even after the PR is closed**.\n\n" - "Thanks for contributing to LiteLLM. We know auto-closures can sting; " - "the goal is to keep the project healthy, not to dismiss your work." - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def close_pr( - pr: dict, - score: int, - threshold: int, - age_days: int, - repo: str | None, - dry_run: bool, - label: str | None, -) -> None: - """Post the explanatory comment and close the PR.""" - pr_number = pr["number"] - repo_args = ["--repo", repo] if repo else [] - - if dry_run: - print( - f" [DRY RUN] Would close PR #{pr_number} " - f"(age={age_days}d, greptile={score}/5): {pr['title']}" - ) - return - - comment_body = format_close_comment(score, threshold) - gh("pr", "comment", str(pr_number), "--body", comment_body, *repo_args) - - if label: - try: - gh("pr", "edit", str(pr_number), "--add-label", label, *repo_args) - except subprocess.CalledProcessError as exc: - stderr = (exc.stderr or "").strip() - print(f" warn: failed to add label '{label}' to #{pr_number}: {stderr}") - - gh("pr", "close", str(pr_number), *repo_args) - print(f" Closed PR #{pr_number} (greptile={score}/5, age={age_days}d)") - - -def evaluate_pr( - pr: dict, - now: dt.datetime, - min_age_days: int, - min_score: int, - repo: str | None, - optout_labels: set[str], - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> tuple[str, int | None, int | None]: - """Decide what to do with `pr` on this triage run. - - Returns (action, score_or_none, age_days_or_none) where action is one of: - "skip-too-young", "skip-optout-label", "skip-not-allowlisted", - "skip-internal", "skip-no-greptile-score", "skip-score-ok", - "warn-grace", "skip-in-grace-period", or "close". - - Drafts are NOT skipped — the goal is "open PR count == PRs internal - collaborators need to action on", and a draft that Greptile scored <4/5 - is still in that queue. Authors can opt out via the `wip` label (see - `DEFAULT_OPTOUT_LABELS`) if they need to keep a long-lived draft open. - - Grace-period semantics: the first time a PR fails the rubric, the - action is `warn-grace` — the caller should post a warning comment but - NOT close the PR. On a subsequent run, if the warning is still less - than `GRACE_PERIOD_SECONDS` old AND the PR still fails, the action is - `skip-in-grace-period`. Once the warning ages out and the rubric is - still failing, the action is `close`. - """ - if has_optout_label(pr, optout_labels): - return ("skip-optout-label", None, None) - - created = parse_iso8601(pr["createdAt"]) - age_days = (now - created).days - # `min_age_days` defaults to 0 (close as soon as Greptile scores low). - # Set a positive value via --min-age-days for one-off backfill runs that - # want to skip very-young PRs. - if min_age_days > 0 and age_days < min_age_days: - return ("skip-too-young", None, age_days) - - # While the allowlist is active it is the sole author gate: only those - # logins are acted on and the external-only restriction is bypassed for - # them. Otherwise auto-close only external OSS contributors — internal - # contributors (BerriAI org members) handle their own backlog. - login = ((pr.get("author") or {}).get("login") or "").lower() - if allowlist: - if login not in allowlist: - return ("skip-not-allowlisted", None, age_days) - elif not is_external_pr_author(pr, repo): - return ("skip-internal", None, age_days) - - comments = fetch_pr_comments(pr["number"], repo) - extraction = extract_greptile_score(comments) - if extraction is None: - return ("skip-no-greptile-score", None, age_days) - - score, _ = extraction - if score >= min_score: - return ("skip-score-ok", score, age_days) - - grace_age = seconds_since_last_grace_warning(comments, now=now) - if grace_age is None: - return ("warn-grace", score, age_days) - if grace_age < GRACE_PERIOD_SECONDS: - return ("skip-in-grace-period", score, age_days) - - return ("close", score, age_days) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--repo", - type=str, - default=None, - help="Repository (owner/repo). Auto-detected if omitted.", - ) - parser.add_argument( - "--min-age-days", - type=int, - default=0, - help=( - "Minimum age (in days) before a PR is eligible. Default 0 = " - "close as soon as Greptile flags it. Set a positive value for " - "one-off backfill runs that want to spare very-young PRs." - ), - ) - parser.add_argument( - "--min-score", - type=int, - default=4, - choices=range(1, 6), - help="Greptile score below which a PR is closed (default: 4 -> closes <4/5).", - ) - parser.add_argument( - "--optout-label", - action="append", - default=None, - help=( - "Label(s) that exempt a PR from auto-close. Repeat to add more. " - "Case-insensitive. When omitted, defaults to " - f"{list(DEFAULT_OPTOUT_LABELS)!r}; passing this flag REPLACES the " - "defaults (argparse `append` with a mutable default would append " - "instead, which we explicitly avoid)." - ), - ) - parser.add_argument( - "--close-label", - type=str, - default=None, - help=( - "Optional label to add to PRs that get auto-closed " - "(e.g. 'auto-closed-low-quality'). Must already exist on the repo." - ), - ) - parser.add_argument( - "--close", - action="store_true", - help="Actually close matching PRs (default is dry-run).", - ) - parser.add_argument( - "--limit", - type=int, - default=None, - help="Maximum number of PRs to close in one run (safety net).", - ) - args = parser.parse_args() - - dry_run = not args.close - if dry_run: - print("=== DRY RUN MODE (pass --close to actually close PRs) ===\n") - - print("Fetching open PRs...") - prs = fetch_open_prs(args.repo) - print(f"Found {len(prs)} open PRs.\n") - - now = dt.datetime.now(dt.timezone.utc) - optout_labels = set(args.optout_label or DEFAULT_OPTOUT_LABELS) - - closed = 0 - summary = { - "close": 0, - "warn-grace": 0, - "skip-in-grace-period": 0, - "skip-too-young": 0, - "skip-optout-label": 0, - "skip-not-allowlisted": 0, - "skip-internal": 0, - "skip-no-greptile-score": 0, - "skip-score-ok": 0, - } - - # `warned` tracks grace-warning comments posted in this run so the - # `--limit` safety net bounds *all* destructive write actions, not - # just closures. Without this cap, a backlog of PRs failing the - # threshold simultaneously could flood contributors with comments. - warned = 0 - for pr in sorted(prs, key=lambda p: p["createdAt"]): - try: - action, score, age_days = evaluate_pr( - pr, - now, - args.min_age_days, - args.min_score, - args.repo, - optout_labels, - ) - summary[action] = summary.get(action, 0) + 1 - - if action == "warn-grace": - assert score is not None - print( - f"#{pr['number']}: \"{pr['title']}\" " - f"(age={age_days}d, greptile={score}/5) -> warn-grace" - ) - post_grace_warning( - pr, - score=score, - threshold=args.min_score, - repo=args.repo, - dry_run=dry_run, - ) - if not dry_run: - warned += 1 - if args.limit is not None and (warned + closed) >= args.limit: - print( - f"\nReached --limit={args.limit} " - f"(closed={closed}, warned={warned}); stopping." - ) - break - continue - - if action != "close": - continue - - assert score is not None and age_days is not None - print( - f"#{pr['number']}: \"{pr['title']}\" " - f"(age={age_days}d, greptile={score}/5) -> close" - ) - close_pr( - pr, - score=score, - threshold=args.min_score, - age_days=age_days, - repo=args.repo, - dry_run=dry_run, - label=args.close_label, - ) - - if not dry_run: - closed += 1 - if args.limit is not None and (warned + closed) >= args.limit: - print( - f"\nReached --limit={args.limit} " - f"(closed={closed}, warned={warned}); stopping." - ) - break - except Exception as exc: # noqa: BLE001 - per-PR errors don't abort the sweep - summary["error"] = summary.get("error", 0) + 1 - print( - f"!! PR #{pr.get('number')}: {exc}", - file=sys.stderr, - ) - continue - - print("\n=== Summary ===") - for key, value in summary.items(): - print(f" {key:28s} {value}") - if dry_run: - print(f"\nTotal would close: {summary['close']}") - else: - print(f"\nTotal closed: {closed}") - print( - f"Total {'would warn (grace)' if dry_run else 'warned (grace)'}: " - f"{summary['warn-grace']}" - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/scripts/triage-requirements.txt b/.github/scripts/triage-requirements.txt deleted file mode 100644 index a18f05fbb95..00000000000 --- a/.github/scripts/triage-requirements.txt +++ /dev/null @@ -1,282 +0,0 @@ -# Hash-pinned dependency set for the Agent Shin triage scripts. -# Installed in privileged triage workflows, so every package is pinned to an -# exact version with SHA-256 hashes and installed with pip --require-hashes. -# -# Regenerate after bumping openai: -# echo 'openai==' \ -# | uv pip compile - --generate-hashes --python-version 3.12 \ -# --no-annotate --no-header -o .github/scripts/triage-requirements.txt - -annotated-types==0.7.0 \ - --hash=sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53 \ - --hash=sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89 -anyio==4.14.0 \ - --hash=sha256:b47c1f9ccf73e67021df785332508f99379c68fa7d0684e8e3492cb1d4b23f89 \ - --hash=sha256:dd9b7a2a9799ed6552fde617b2c5df02b7fdd7d88392fc48101e51bae46164d9 -certifi==2026.6.17 \ - --hash=sha256:024c88eeec92ca068db80f02b8b07c9cef7b9fe261d1d535abfd5abd6f6af432 \ - --hash=sha256:2227dcbaafe0d2f59279d1762ddddc37783ed4354594f194ffc31d20f41fc3db -distro==1.9.0 \ - --hash=sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed \ - --hash=sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2 -h11==0.16.0 \ - --hash=sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1 \ - --hash=sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86 -httpcore==1.0.9 \ - --hash=sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55 \ - --hash=sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8 -httpx==0.28.1 \ - --hash=sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc \ - --hash=sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad -idna==3.18 \ - --hash=sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2 \ - --hash=sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848 -jiter==0.15.0 \ - --hash=sha256:01a8222cf05ab1128e239421156c207949808acaaea2bdfd33130ae666786e86 \ - --hash=sha256:032396229564bca02440396bd327710719f724f5e7b7e9f7a8eb3faa4a2c2281 \ - --hash=sha256:04b400bbf8c9efb03d9bdd976475c919c1d85593b04b9fff7ae234065daf87ae \ - --hash=sha256:05906b93d72f03339e6bb7cf8dc10ebda64a0266126eed6beba79e20abcf5fd4 \ - --hash=sha256:066f8f33f18b2419cd8213b2436fa7fbc9c499f315971cfa3ce1f9820c001b1b \ - --hash=sha256:0ab068bce62a45aa3e7367eceaffb5dde60b7eb853be8dece45132e3d0ff4879 \ - --hash=sha256:0be6f5ad41a809f303f416d17cec92a7a725902fb9b4f3de3d19362ac0ef8554 \ - --hash=sha256:0e90a1c315a0226ec822d973817967f9223b7701546c8c2a7913e7ab0926294d \ - --hash=sha256:0f862193b8696249d22ec433e85fd2ab0ad9596bc3e45e6c0bc55e8aeba97be2 \ - --hash=sha256:1303d4d68a9b051ea90502402063ecf3807da00ad2affa19ca1ae3b90b3c5f67 \ - --hash=sha256:144f8e72cb53dab146347b91cceac01f5481237f2b93b4a339a1ee8f8878b67c \ - --hash=sha256:182226cbc930c9fab81bc2e41a4da672f89539906dadb05e75670ac07b94f71f \ - --hash=sha256:1c11465f97e2abf45a014b83b730222f8f1c5335e802c7055a67d50de6f1f4e3 \ - --hash=sha256:1c15024a3d892223b18f597c86d59387249dc396590844ce6b9f6131d1093bae \ - --hash=sha256:1d54fb5b31dea401a41af3f8a7d2512e9b6a6a005491e6166c7e4ffab9639a9c \ - --hash=sha256:25ffbe229aa8cd98c28879d8aa1a6e34ae77992ab984a65fba800859dab16269 \ - --hash=sha256:2a77aadd57cac1682e4401a72724d2796d89a4ba129b1a5812aa94ee480826eb \ - --hash=sha256:2ae901f3a55bfafdde31d289590fa25e3245735a2b1e8c7cc15871710a002871 \ - --hash=sha256:2b0074e2f56eb2dacca1689760fd2852a068f85a0547a157b82cb4cafeb6768b \ - --hash=sha256:2c8aea7781d2a372227871de4e1a1332aa96f5a89fd76c5e835dafdbad102887 \ - --hash=sha256:2c9cb907439d20bd0c7d7565ca01ee52234203208433749bae5b516907526928 \ - --hash=sha256:2fb6a5d26af81fc0f00f9360a891e05cf755e149bba391c4d563adc54812973d \ - --hash=sha256:2fd73e3da91a0a722d67165e849ce2cdc10de0e0d48738c142be8c6c5f310f4c \ - --hash=sha256:30ce1a5d16b5641dc935d50ef775af6a0871e3d14ab05d6fc54dff371b78e558 \ - --hash=sha256:30ce785d2adb8e32c3f7741442370a74834ec4c01f3c48f0750227a0b4ef27d6 \ - --hash=sha256:30f2218e6a9e5c18bc10fe6d41ac189c442c88eacf11bad9f28ef95a9bef00e6 \ - --hash=sha256:351a341c2105aa430b7047e30f1bf7975f6313b00165d3fc07be2edaf741f279 \ - --hash=sha256:37a10c377ce3a4a85f4a67f28b7afe093154cde77eaf248a72e856aa08b4d865 \ - --hash=sha256:392b8ab019e5502d08aff85c6272209c24bc2cbe706ea82a56368f524236614a \ - --hash=sha256:3e4540b8e74e4268811ac05db226a6a128ff572e7e0ce3f1163b693cadb184cd \ - --hash=sha256:40b2c7e92c44a84d748d21706c68dc6ff8161d80b59c99d774721a0d2317d7c7 \ - --hash=sha256:411fa4dfa5a7ae3d11491027ffb9beadec3996010a986862db70d91abba1c750 \ - --hash=sha256:4251acc80e2b7c9b7b8823456ea0fceeb0734dac2df7636d3c711b38476b5a76 \ - --hash=sha256:42bfb257930800cf43e7c62c832402c704ab60797c992faf88d20e903eac8f32 \ - --hash=sha256:4363818355dbc70ae1a8e9eaba9de350d93ede4ff6992b8f8eb8cbb6e5122d42 \ - --hash=sha256:4ab395feec8d249ec4044e228e98a7033f043426a265df439dc3698823f0a4e4 \ - --hash=sha256:50164d7610c00e7cd913a873fce30b6beeebf4b37e53983e33f22de4c900f6b8 \ - --hash=sha256:50e51156192722a9c58db112837d3f8ef96fb3c5ecc14e95f409134b08b158ec \ - --hash=sha256:510c8b3c17a0ed9ac69850c0438dada3c9b82d9c4d589fcb62002a5a9cf3a866 \ - --hash=sha256:5157de9f76eb4bc5ea74a1219366a25f945ad305641d74e04f59c54087091aa9 \ - --hash=sha256:54d5d6090cdc1b7c9e780dfb04949a990adb1e301a2fc0bbcee7de4638d33f9a \ - --hash=sha256:553fcac2ef2cb990877f9fc0833b8b629a3e6a5670b6b5fd58219b41a653ddc4 \ - --hash=sha256:5607e6013ed7e6b0ec9661e467b7ffde0aa7ab36833a04850f26fcf88ed4845b \ - --hash=sha256:5d6a60072b44c3c2b797a7ddcbcbbf2b34ea3cfd4721580fbfd2a09d9d9b84ba \ - --hash=sha256:5f30bae8bc1c2d613e28e5af3e8cceb09b742f1c8a8a5f839fb67afaffc03b61 \ - --hash=sha256:62ebd14e47e9aed9df4472afcb2663668ce4d74891cd54f86bf6e44029d6dc89 \ - --hash=sha256:631f13a3d04e97d4e083993b10f4b99530e3a10d953e2eb5e196b7dc7f812ce0 \ - --hash=sha256:6550fa135c7deb8ead6af49ed7ff648532ea8334a1447fe34a36315ef79c5c29 \ - --hash=sha256:66b1880df2d01e206e8339769d1c7c1753bcb653efd6289e203f6f24ebada0c0 \ - --hash=sha256:6eac374c5c975709b69c10f09afd199df74150172156ad10c8d4fd785b7da995 \ - --hash=sha256:71683c38c825452999b5717fcae07ea708e8c93003e808be4319c1b02e3d176e \ - --hash=sha256:7553333dd0930c104a5a0db8df72bf7219fe663d731383b576bb6ed6351c984d \ - --hash=sha256:75e8a04e91432dde9f1838373cf93d23726c79d3e908d319acf0e796f85592e7 \ - --hash=sha256:773b6eb282ce11ee19f05f6b2d4404fa308e5bbd353b0b80a0262caad6db2cd7 \ - --hash=sha256:774f93f65031856bf14ad9f59bdcab8b8cad501e5ceabd51ba3525f76937a25b \ - --hash=sha256:7c468136b8bd6bb18c8786e4236a1fa27362f24cb23450ba0cb204ab379b8e6f \ - --hash=sha256:7ce8902f939970048b233087082e7bb829db29375811c7ad50687b8624c6fd08 \ - --hash=sha256:7d3d6683288c11cbab50e865f2e2f13950179aa45410e30b2cfbd3fb7b0177bf \ - --hash=sha256:7f6163c0f10b055245f814dcc59f4818da60dfe72f3e72ab89fc24b6bd5e9c52 \ - --hash=sha256:8020c99ec13a7db2b6f96cbe82ef4721c88b426a4892f27478044af0284615ef \ - --hash=sha256:813dfbb17d65328bf86e5f0905dd277ba2265d3ca20556e86c0c7035b7182e5a \ - --hash=sha256:860a74063284a2ae9bfedd694f299cc2c68e2696c5f3d440cc9d18bb81b9dd04 \ - --hash=sha256:8c9004af7c8d67cce7f1aae1026fb55607f4aa600710d08ede3a3ce4aeefe7e0 \ - --hash=sha256:8d2c0c44d569ce0f2850f5c926f8caeb5f245fbc84475aeb36efccc2103e6dbd \ - --hash=sha256:8f7e9bc0f1135039b22ee6eab588d42df1ce55842b30740a352885eb267bd941 \ - --hash=sha256:90c5db5527c221249a876160663ab891ace358c17f7b9c93ec1478b7f0550e5c \ - --hash=sha256:9100ddbec09741cc66feb0fc6773f8bdbd0e3c345689368f260082ff85dcc0cd \ - --hash=sha256:913d02d29c9606643418d9ccfc3b72492ab25a6bf7889934e09a3490f8d3438b \ - --hash=sha256:980c256edb05b78a111b99c4de3b1d32e31634b867fd1fc2cf726e7b7bba9854 \ - --hash=sha256:9f924585cdacf631cd382b657966847bb537bf9ed0a6f9b991da5f05a631480f \ - --hash=sha256:a254e10b593624d230c365b6d616b22ca0ad65e63a16e6631c2b3466022e6ba8 \ - --hash=sha256:a2a438005b6f22d0273413484d6094d7c2c5d10ec1b3a3bf128e0d1d3ba53258 \ - --hash=sha256:a97261f1fccb8e50ecd2890a96e46efdc3f57c80a197324c6777827231eca712 \ - --hash=sha256:ab596fa3837e91e7e6a31b5f639988bfc6a35d1f915ac3932d946062219d588f \ - --hash=sha256:abbf258599526ad0326fe51e252e24f2bd6f24f1852681b4b78feda3808f1d18 \ - --hash=sha256:ac0d9ddea4350974be7a221fc25895f251a8fee748c889bdced2141c0fec1a49 \ - --hash=sha256:acf4ee4d1fc55917239fe72972fb292dd773055d05eb040d36f4326e02cc2c0e \ - --hash=sha256:ae1b0d82ac2d987f9ea512b1c9adfcc71a28de3dea3a6039b54d76cffda9901e \ - --hash=sha256:b15741f501469009ae0ae90b7147958a664a7dede40aa7ff174a8a4645f546d0 \ - --hash=sha256:b15d3ec9b0449c40e85319bdb4caa8b77ab526e74f5532ed94bec15e2f66822c \ - --hash=sha256:b3b3b775e33d3bfaec9899edc526ae97b0da0bf9d071a46124ba419149a414f8 \ - --hash=sha256:b6c0ffae686c39bf3737be60793783267628783ea42545632c10b291105aee45 \ - --hash=sha256:c210f8b35dc6f30aafd4b4365ca89b9d1189f21ab49b8e68fa6322a847aef138 \ - --hash=sha256:c2f6bb8b5216ab9e7873bc08b5d7bef2b8abbb578a3069bf1cd14a45d71d771d \ - --hash=sha256:c60e71b6d10cfc284c9bf36bd885e8d44c46f688ce50aa91b5edd90181dea687 \ - --hash=sha256:c6694a173ecabc12eb60efbc0b474464ead1951ff65cd8b1e72100715c64512b \ - --hash=sha256:c77496cb10bd7549690fbbab3e5ec05857b83e49276f4a9423a766ddd2afcd4c \ - --hash=sha256:c84c1b7be454b0c16f8499b4ebfbfd82ea5cca6527cceefcbbc06a7557b5ed2e \ - --hash=sha256:cc0bc345cf2df9d1c00ac443f50d543c1ccfa8b0422cb85b1ab70d681c0b255b \ - --hash=sha256:ceb8fc27d38793f9c97149be8302720c5b22e5c195a37bf2c45dc36c4600a512 \ - --hash=sha256:cf4bd113a69c0a740e27cb962ce10630c36d2b8f59d759a651b955ee9d18a823 \ - --hash=sha256:d1aa62e277fc1cbd80e6deacae6f4d983b41b3d7728e0645c5d741a6149bba45 \ - --hash=sha256:d1e7b1776f0797956c509e123d0952d10d293a9492dea9f288ab9570ec01d1a5 \ - --hash=sha256:d636d5095155afd364247f65070fab7beda13498d7ff4de331046e704ab9657f \ - --hash=sha256:d726e3ceeb337191324b49de298142f27c3ad10886341555d1d5315b5f252c6a \ - --hash=sha256:d72d8af5c1013656a8870c866660627d1a75bc185814ee022c8533caa1de88ae \ - --hash=sha256:d8d2955167274e15d79a7a020afdd9b39c990eb80b2d89fca695d92dcfdd38ec \ - --hash=sha256:d92a5cd21fdb083931d546c207aa29633787c5dc5b02daab2d32b843f88a2c53 \ - --hash=sha256:e58585a58209d72691ce2d62a9147445f5a87beb0bde97fde284c96ae392a3d1 \ - --hash=sha256:e7196e56f1cd69af1dbb07dff02dcfb260a50b45a82d409d92a06fedb32473b5 \ - --hash=sha256:eda3071db3346334beae1360b46da4606da57bf3528c167b3c38533afaf9f2c5 \ - --hash=sha256:edebcf7d1f601199084bb6e844d7dc67e03e04f6ac786b0332d616635c4ff7a4 \ - --hash=sha256:ef1fd24d9413f6209e00d3d5a453e67acfe004a25cc6c8e8484faed4311ab9e8 \ - --hash=sha256:f0b271b462769543716f92d3a4f90527df6ef5ed05ee95ec4137f513e21e1b77 \ - --hash=sha256:f18f85e4218d1b40f000f42a92239a7a61a902cd42c65e6c360dbd17dcb20894 \ - --hash=sha256:f1e1754960f38ec40613a07e5e372df67acb3b890fb383b6fb3de3e49ddbf3c7 \ - --hash=sha256:f2143ab06181d2b029eedcb6af3cebe95f11bbac62441781860f98ee9330a6a6 \ - --hash=sha256:f3d37768fce7f88dd2a8c6091f2325dea27d30d30d5c6e7a1c0f0af77723b708 \ - --hash=sha256:fa248c9eb220197d363f688818dac2fd4b2f0cd7d843ca7105d652034823427d -openai==2.33.0 \ - --hash=sha256:03ac37d70e8c9e3a8124214e3afa785e2cbc12e627fbd98177a086ef2fd87ad5 \ - --hash=sha256:f850c435e2a4685bba3295bd54912dd26315d9c1b7733068186134d6e0599f9a -pydantic==2.13.4 \ - --hash=sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba \ - --hash=sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6 -pydantic-core==2.46.4 \ - --hash=sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0 \ - --hash=sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262 \ - --hash=sha256:041bde0a48fd37cf71cab1c9d56d3e8625a3793fef1f7dd232b3ff37e978ecda \ - --hash=sha256:0c563b08bca408dc7f65f700633d8442fffb2421fc47b8101377e9fd65051ff0 \ - --hash=sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e \ - --hash=sha256:0ce40cd7b21210e99342afafbd4d0f76d784eb5b1d60f3bdc566be4983c6c73b \ - --hash=sha256:0e96592440881c74a213e5ad528e2b24d3d4f940de2766bed9010ab1d9e51594 \ - --hash=sha256:10e17cbb10a330363733efc4d7c4d0dd827ac0909b8f6a6542298fed1ea62f29 \ - --hash=sha256:133878133d271ade3d41d1bfb2a45ec38dbdbda40bc065921c6b04e4630127e2 \ - --hash=sha256:14d4edf427bdcf950a8a02d7cb44a08614388dd6e1bdcbf4f67504fa7887da9c \ - --hash=sha256:14f4c5d6db102bd796a627bbb3a17b4cf4574b9ae861d8b7c9a9661c6dd3362d \ - --hash=sha256:17299feefe090f2caa5b8e37222bb5f663e4935a8bfa6931d4102e5df1a9f398 \ - --hash=sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d \ - --hash=sha256:18e5ceec2ab67e6d5f1a9085e5a24c9c4e2ac4545730bfe668680bca05e555f3 \ - --hash=sha256:19e51f073cd3df251856a8a4189fbdf1de4012c3ebacfb1884f94f1eb406079f \ - --hash=sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb \ - --hash=sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7 \ - --hash=sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5 \ - --hash=sha256:228ee9bae8bef5b1e97ec58302f80357c37199e0d0a99174e138d28e6957b9d9 \ - --hash=sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462 \ - --hash=sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4 \ - --hash=sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b \ - --hash=sha256:2f84c03c8607173d16b5a854ec68a2f9079ae03237a54fb506d13af47e1d018d \ - --hash=sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df \ - --hash=sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2 \ - --hash=sha256:3447661d99f75a3683a4cf5c87da72f2161964611864dbbeac7fbb118bb4bfc0 \ - --hash=sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519 \ - --hash=sha256:395aebd9183f9d112f569aeb5b2214d1a10a33bec8456447f7fbdfa51d38d4cd \ - --hash=sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7 \ - --hash=sha256:3be77f45df024d789a672ae34f8b06fb346c4f9f46ea714956660ea4862e89ac \ - --hash=sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6 \ - --hash=sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565 \ - --hash=sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898 \ - --hash=sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb \ - --hash=sha256:432c179df7874eeb73307aad2df0755e1ae0efa61ff0ea89b93e194411ae3928 \ - --hash=sha256:4a05d69cba51d852c5c3e92758653245a50c0b646ced0cf05bd793ed592839d6 \ - --hash=sha256:4c63ebc82684aa89d9a3bcbd13d515b3be44250dc68dd3bd81526c1cb31286c3 \ - --hash=sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a \ - --hash=sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596 \ - --hash=sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987 \ - --hash=sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e \ - --hash=sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d \ - --hash=sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712 \ - --hash=sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008 \ - --hash=sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd \ - --hash=sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1 \ - --hash=sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be \ - --hash=sha256:66ce7632c22d837c95301830e111ad0128a32b8207533b60896a96c4915192ea \ - --hash=sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292 \ - --hash=sha256:6f2eeda33a839975441c86a4119e1383c50b47faf0cbb5176985565c6bb02c33 \ - --hash=sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3 \ - --hash=sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4 \ - --hash=sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b \ - --hash=sha256:7bfb192b3f4b9e8a89b6277b6ce787564f62cfd272055f6e685726b111dc7826 \ - --hash=sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac \ - --hash=sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7 \ - --hash=sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d \ - --hash=sha256:8358a950c8909158e3df31538a7e4edc2d7265a7c54b47f0864d9e5bae9dcebf \ - --hash=sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4 \ - --hash=sha256:86e1a4418c6cd97d60c95c71164158eaf7324fae7b0923264016baa993eba6fc \ - --hash=sha256:8b9bab013d1c7a79d3501ff86d0bc9c31bf587db4551677b96bec07df78c6b15 \ - --hash=sha256:8c5dac79fa1614d1e06ca695109c6105923bd9c7d1d6c918d4e637b7e6b32fd3 \ - --hash=sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b \ - --hash=sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914 \ - --hash=sha256:9037063db01f09b09e237c282b6792bd4da634b5402c4e7f0c61effed7701a04 \ - --hash=sha256:905a0ed8ea6f2d61c1738835f99b699348d7857379083e5fc497fa0c967a407c \ - --hash=sha256:90884113d8b48f760e9587002789ddd741e76ab9f89518cd1e43b1f1a52ec44b \ - --hash=sha256:91a06d2e259ecfbd8c901d70c3c507900458498142b3026a296b7de4d1322cc9 \ - --hash=sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce \ - --hash=sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4 \ - --hash=sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a \ - --hash=sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f \ - --hash=sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424 \ - --hash=sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894 \ - --hash=sha256:9bc519fbf2b7578398853d815009ae5e4d4603d12f4e3f91da8c06852d3da3e9 \ - --hash=sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76 \ - --hash=sha256:9f444c499b3eefd3a92e348059471ea0c3a6e303d9c1cec09fa748fd9f895201 \ - --hash=sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb \ - --hash=sha256:a0f62d0a58f4e7da165457e995725421e0064f2255d8eccebc49f41bbc23b109 \ - --hash=sha256:a396dcc17e5a0b164dbe026896245a4fa9ff402edca1dff0be3d53a517f74de4 \ - --hash=sha256:aaa2a54443eff1950ba5ddc6b6ccda0d9c84a364276a62f969bdf2a390650848 \ - --hash=sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526 \ - --hash=sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0 \ - --hash=sha256:b078afbc25f3a1436c7a1d2cd3e322497ee99615ba97c563566fdf46aff1ee01 \ - --hash=sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458 \ - --hash=sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e \ - --hash=sha256:bb63e0198ca18aad131c089b9204c23079c3afa95487e561f4c522d519e55aba \ - --hash=sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a \ - --hash=sha256:c1747f85cee84c26985853c6f3d9bd3e75da5212912443fa111c113b9c246f39 \ - --hash=sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c \ - --hash=sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000 \ - --hash=sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b \ - --hash=sha256:c7a7bd4e39e8e4c12c39cd480356842b6a8a06e41b23a55a5e3e191718838ddf \ - --hash=sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4 \ - --hash=sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd \ - --hash=sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28 \ - --hash=sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9 \ - --hash=sha256:d396ec2b979760aaf3218e76c24e65bd0aca24983298653b3a9d7a45f9e47b30 \ - --hash=sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983 \ - --hash=sha256:d80ee3d731373b24cebbc10d689ca4ee1875caf0d5703a245db18efd4dd37fc1 \ - --hash=sha256:d995260fdf4e1db774581b4900e0f832abe3c7c84996726bbc161b19c8f29e76 \ - --hash=sha256:da4b951fe36dc7c3a1ccb4e3cd1747c3542b8c9ceede8fc86cae054e764485f5 \ - --hash=sha256:daa27d92c36f24388fe3ad306b174781c747627f134452e4f128ea00ce1fe8c4 \ - --hash=sha256:db06ffe51636ffe9ca531fe9023dd64bdd794be8754cb5df57c5498ae5b518a7 \ - --hash=sha256:e0d65b8c354be7fb5f720c3caa8bc940bc2d20ce749c8e06135f07f8ed95dd7c \ - --hash=sha256:e68b7a074f65a2fd746c52a7ce6142ab7006074ac269ace0c25cd8ba171f8066 \ - --hash=sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3 \ - --hash=sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02 \ - --hash=sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89 \ - --hash=sha256:ea793e075b70290d89d8142074262885d3f7da19634845135751bd6344f73b50 \ - --hash=sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76 \ - --hash=sha256:f13a646d65d09fbf1bc6b3a9635d30095c8e7e5cc419ff35ecc563c5fd04cd49 \ - --hash=sha256:f47286a97f0bc9b8859519809077b91b2cefe4ae47fcbf5e466a009c1c5d742b \ - --hash=sha256:f747929cf940cddb5b3668a390056ddd5ba2e5010615ea2dcf4f9c4f3ab8791d \ - --hash=sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7 \ - --hash=sha256:f9fa868638bf362d3d138ea55829cefb3d5f4b0d7f142234382a15e2485dbec4 \ - --hash=sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c \ - --hash=sha256:fc010ab034c8c7452522748bf937df58020d256ccae0874463d1f4d01758af8e \ - --hash=sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff \ - --hash=sha256:fd8b3d9fd264be37976686c7f65cd52a83f5e84f4bfd2adf9c1d469676bbb6ae -sniffio==1.3.1 \ - --hash=sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2 \ - --hash=sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc -tqdm==4.68.3 \ - --hash=sha256:00dfa48452b6b6cfae3dd9885636c23d3422d1ec97c66d96818cbd5e0821d482 \ - --hash=sha256:39832cc2def2789a6f29df83f172db7416cea70052c0907a57801c5f2fdccb03 -typing-extensions==4.15.0 \ - --hash=sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466 \ - --hash=sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548 -typing-inspection==0.4.2 \ - --hash=sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7 \ - --hash=sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464 diff --git a/.github/scripts/triage_with_llm.py b/.github/scripts/triage_with_llm.py deleted file mode 100644 index e23a012425a..00000000000 --- a/.github/scripts/triage_with_llm.py +++ /dev/null @@ -1,1797 +0,0 @@ -#!/usr/bin/env python3 -""" -Agent Shin — LLM-as-judge triage for external OSS pull requests and issues. - -Evaluates a single PR or issue against the contribution rubric and, when the -LLM judge marks it as failing, posts an explanatory comment + closes the -PR/issue. Re-triggers on `reopened` so contributors can iterate back in by -filling in the missing pieces and reopening. - -Internal BerriAI contributors (`author_association` in {OWNER, MEMBER, -COLLABORATOR}) and bot accounts are skipped entirely. - -Usage: - triage_with_llm.py --repo owner/repo --pr 1234 - triage_with_llm.py --repo owner/repo --issue 5678 - triage_with_llm.py --repo owner/repo --pr 1234 --close # actually close - triage_with_llm.py --repo owner/repo --pr 1234 --print-prompt # show prompt - -Defaults are SAFE: without `--close` the script writes a verdict to stdout (and, -when running in GitHub Actions, to $GITHUB_STEP_SUMMARY) but takes no GitHub -write actions. - -Environment: - GH_TOKEN / GITHUB_TOKEN - for `gh` CLI auth (auto-set in Actions) - OPENAI_API_KEY - required when --close is passed - OPENAI_BASE_URL - optional (route to any OpenAI-compatible API) - TRIAGE_MODEL - optional model override (default: gpt-5.4-mini) -""" - -from __future__ import annotations - -import argparse -import datetime as dt -import json -import os -import re -import subprocess -import sys -import textwrap -import urllib.parse -from typing import Any, Iterable - -# Add this script's directory to `sys.path` so the sibling -# `agent_shin_shared` module is importable when the script is invoked -# directly (e.g. `python3 .github/scripts/triage_with_llm.py ...`) and -# also when the tests load this script via -# `importlib.util.spec_from_file_location`. -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from agent_shin_shared import ( # noqa: E402 -- sys.path adjusted above - AGENT_SHIN_CLOSE_MARKER, - AGENT_SHIN_DEFAULT_BOT_LOGIN, - ALLOWLIST_LOGINS, - GRACE_COMMENT_MARKER, - GRACE_PERIOD_SECONDS, - GREPTILE_BOT_LOGINS, - SCORE_PATTERN, - extract_greptile_score, - gh, - parse_iso8601, - seconds_since_latest_marker_comment, -) - -DEFAULT_MODEL = "gpt-5.4-mini" - -INTERNAL_ASSOCIATIONS = frozenset({"OWNER", "MEMBER", "COLLABORATOR"}) - -# `AGENT_SHIN_DEFAULT_BOT_LOGIN` is imported from `agent_shin_shared`. -# When the workflow uses the default `secrets.GITHUB_TOKEN`, the -# closure / reopen event's `actor.login` is `github-actions[bot]`. The -# env override `AGENT_SHIN_BOT_LOGIN` exists for local debugging and for -# repos that wire Agent Shin to a PAT. - -# HTML marker appended to every reconsider verdict comment. We grep for this -# on subsequent reconsider triggers to enforce a short cooldown so that -# repeated `@agent-shin reconsider` comments don't burn CI/LLM budget. -# Using a unique HTML comment keeps the marker invisible to humans while -# being trivially greppable from a comments-list API response. -RECONSIDER_COMMENT_MARKER = "" - -# Minimum gap between two reconsider verdicts on the same PR/issue. Set to -# 10 minutes — long enough that a contributor can't trivially spam the -# trigger, short enough that a genuine "I just pushed a fix and reupdated -# the body" iteration loop isn't punished. -RECONSIDER_RATE_LIMIT_SECONDS = 600 - -# `GRACE_COMMENT_MARKER` (HTML marker on the grace-period warning comment -# posted on the first low-quality detection — used on subsequent triage -# runs to detect that a warning was already posted and measure how long -# ago it was posted) and `GRACE_PERIOD_SECONDS` (length of the grace -# period between the warning and the actual auto-close, 2 hours) are -# imported from `agent_shin_shared` so the daily Greptile sweep and the -# LLM judge agree on the same marker and duration. - -# --- Review-gate ("ready for review" label lifecycle) configuration ---------- -# The review gate keeps a single label in sync with whether a PR currently -# clears BOTH quality bars: the LLM rubric (clear problem + expected/actual + -# QA proof, or a linked issue) AND Greptile's most recent confidence score. -READY_FOR_REVIEW_LABEL = "ready for review" -DEFAULT_GRACE_DAYS = 1 # 24h before an un-passing, un-tagged PR is auto-closed -DEFAULT_MIN_GREPTILE_SCORE = 4 # Greptile < 4/5 counts as "not passing" - -# Hidden HTML-comment markers stamped into review-gate comments. They never -# render in the GitHub UI but let the gate detect its own prior actions so it -# (a) posts the within-grace "what's missing" notice at most once and (b) can -# tell a first-time pass ("ready for review") from a recovery after a -# regression ("all clear again"). -READY_MARKER = "" -REGRESSED_MARKER = "" -WITHIN_GRACE_MARKER = "" - -# `GREPTILE_BOT_LOGINS` (Greptile's GitHub App login variants — -# `greptile-apps[bot]` in REST API comments, `greptile-apps` in -# `gh pr view --json` output) and `SCORE_PATTERN` (regex matching lines -# like `Confidence Score: 3/5`) are imported from `agent_shin_shared` -# so the daily sweep and the review gate read the score through the -# same set of logins / patterns. - -# `AGENT_SHIN_CLOSE_MARKER` is imported from `agent_shin_shared` so this LLM -# judge and the daily Greptile sweep stamp the same marker on their close -# comments — `was_closed_by_agent_shin` keys the reconsider reopen path off it. - -# Model families that require `reasoning_effort` to be set, and that reject -# `temperature != 1` unless `reasoning_effort` is "none". For these models we -# pass `reasoning_effort="none"` so a `temperature=0` deterministic judgment -# is still accepted. See litellm/llms/openai/chat/gpt_5_transformation.py for -# the full set of constraints LiteLLM applies to these models. -GPT5_FAMILY_PREFIX = "gpt-5" - -# Regexes for picking off "obvious passes" without burning LLM tokens. -# -# Keep this list to GitHub's documented PR-closing keywords only -# (https://docs.github.com/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue). -# Casual mentions like "see #1234" or "ref #1234" are intentionally NOT -# auto-passed — they should fall through to the LLM judge, which has the -# stricter rubric "a bare issue number without a closing keyword counts only -# if it's clearly the related issue (not a passing mention)". -LINKED_ISSUE_PATTERN = re.compile( - r"\b(?:fixes|fix|fixed|closes|close|closed|resolves|resolve|resolved)\s+" - r"(?:#\d+|https?://github\.com/[\w.-]+/[\w.-]+/issues/\d+)", - re.IGNORECASE, -) -HTML_COMMENT_PATTERN = re.compile(r"", re.DOTALL) - - -# --------------------------------------------------------------------------- -# gh helpers -# -# `gh` is imported from `agent_shin_shared` so a future change (timeout, -# logging, retry) only needs to be made once. - - -def fetch_pr(repo: str, number: int) -> dict: - """Return the full GitHub REST representation of a PR.""" - return json.loads(gh("api", f"repos/{repo}/pulls/{number}")) - - -def fetch_issue(repo: str, number: int) -> dict: - """Return the full GitHub REST representation of an issue.""" - return json.loads(gh("api", f"repos/{repo}/issues/{number}")) - - -def post_comment(repo: str, number: int, body: str) -> None: - """Post an issue-style comment (works for both issues and PRs).""" - gh( - "api", - f"repos/{repo}/issues/{number}/comments", - "-X", - "POST", - "-f", - f"body={body}", - ) - - -def close_pr(repo: str, number: int) -> None: - """Close a pull request (state=closed).""" - gh( - "api", - f"repos/{repo}/pulls/{number}", - "-X", - "PATCH", - "-f", - "state=closed", - ) - - -def reopen_pr(repo: str, number: int) -> None: - """Reopen a previously-closed pull request (state=open). - - Used by the `@agent-shin reconsider` comment-trigger flow: the bot has - write access via GH_TOKEN, so it can reopen on the contributor's behalf - even though GitHub doesn't let the OSS author do it themselves. - """ - gh( - "api", - f"repos/{repo}/pulls/{number}", - "-X", - "PATCH", - "-f", - "state=open", - ) - - -def close_issue(repo: str, number: int, *, not_planned: bool = True) -> None: - """Close an issue, marking state_reason=not_planned by default.""" - args = [ - "api", - f"repos/{repo}/issues/{number}", - "-X", - "PATCH", - "-f", - "state=closed", - ] - if not_planned: - args.extend(["-f", "state_reason=not_planned"]) - gh(*args) - - -def reopen_issue(repo: str, number: int) -> None: - """Reopen a previously-closed issue (state=open, state_reason=reopened).""" - gh( - "api", - f"repos/{repo}/issues/{number}", - "-X", - "PATCH", - "-f", - "state=open", - "-f", - "state_reason=reopened", - ) - - -def add_label(repo: str, number: int, label: str) -> None: - """Add a label to a PR/issue (GitHub creates the label if it's missing).""" - gh( - "api", - f"repos/{repo}/issues/{number}/labels", - "-X", - "POST", - "-f", - f"labels[]={label}", - ) - - -def remove_label(repo: str, number: int, label: str) -> None: - """Remove a label from a PR/issue. A missing label (404) is not an error.""" - encoded = urllib.parse.quote(label, safe="") - try: - gh( - "api", - f"repos/{repo}/issues/{number}/labels/{encoded}", - "-X", - "DELETE", - ) - except subprocess.CalledProcessError as exc: - stderr = (exc.stderr or "").lower() - if "404" in stderr or "not found" in stderr: - return - raise - - -def _iter_paginated_json(*api_args: str) -> Any: - """Yield JSON objects from `gh api --paginate ... -q '.[]'`. - - `gh api --paginate` on a JSON-array endpoint concatenates pages into - one stream; `-q '.[]'` flattens that stream into newline-delimited - objects (jq-style). This keeps memory bounded for chatty endpoints - like issue events/comments on long-lived PRs. - """ - raw = gh("api", "--paginate", *api_args, "-q", ".[]") - for line in raw.splitlines(): - line = line.strip() - if not line: - continue - try: - yield json.loads(line) - except json.JSONDecodeError: - # A malformed line should not blow up the whole guard. Skip and - # carry on — at worst the guard fail-closes (returns False / - # None) and the caller treats it as "unknown". - continue - - -def fetch_last_close_event( - repo: str, number: int -) -> tuple[str | None, dt.datetime | None]: - """Return the actor login and timestamp of the most recent `closed` event. - - Either field may be None: actor when the events API returns nothing - (unusual for a closed item, but possible on transient errors), and - timestamp when the event lacks `created_at` or the value can't be - parsed. `was_closed_by_agent_shin` fail-closes on either. - """ - actor: str | None = None - closed_at: dt.datetime | None = None - for event in _iter_paginated_json(f"repos/{repo}/issues/{number}/events"): - if event.get("event") != "closed": - continue - actor = (event.get("actor") or {}).get("login") - created = event.get("created_at") - if not created: - closed_at = None - continue - try: - closed_at = parse_iso8601(created) - except ValueError: - closed_at = None - return actor, closed_at - - -# How much older than the latest `closed` event the Agent Shin marker -# comment is allowed to be while still counting as "this close was Agent -# Shin's". Agent Shin posts the close comment immediately before closing, -# so the marker timestamp is normally at most a few seconds before the -# close event; the buffer just absorbs clock skew between the comments -# API and the events API. -AGENT_SHIN_CLOSE_MARKER_SKEW_SECONDS = 300 - - -def was_closed_by_agent_shin( - repo: str, number: int, *, bot_login: str | None = None -) -> bool: - """Return True iff Agent Shin itself most-recently closed this PR/issue. - - This is the guard that stops `@agent-shin reconsider` from reopening an - item Agent Shin did not close — a maintainer closing for non-rubric - reasons (security, duplicate, design rejection), or a different workflow - (stale/duplicate sweeps) closing under the shared `github-actions[bot]` - identity. Three independent signals must all hold, because that identity - is not unique to Agent Shin and a marker comment from a prior - closed/reopened cycle would otherwise vouch for an unrelated close: - - 1. The most recent `closed` event's actor is the bot identity. - 2. Agent Shin left one of its auto-close comments, detected via - `AGENT_SHIN_CLOSE_MARKER`. The actor check alone can't tell an - Agent Shin close from any other `github-actions[bot]` close. - 3. That marker comment was posted at (or just before) the latest - close event, not on a previous close in an - Agent-Shin-close -> reconsider-reopen -> other-bot-reclose cycle. - - The check is intentionally fail-closed: any uncertainty about who closed - the item is treated as "not Agent Shin" so the destructive reopen path - stays gated. - """ - expected = ( - bot_login - or os.environ.get("AGENT_SHIN_BOT_LOGIN") - or AGENT_SHIN_DEFAULT_BOT_LOGIN - ).lower() - actor, closed_at = fetch_last_close_event(repo, number) - if not actor or actor.lower() != expected or closed_at is None: - return False - marker_seconds = seconds_since_last_agent_shin_close( - repo, number, bot_login=bot_login - ) - if marker_seconds is None: - return False - close_age_seconds = (dt.datetime.now(dt.timezone.utc) - closed_at).total_seconds() - return marker_seconds <= close_age_seconds + AGENT_SHIN_CLOSE_MARKER_SKEW_SECONDS - - -def _seconds_since_latest_marker_comment( - repo: str, - number: int, - *, - marker: str, - bot_login: str | None = None, -) -> float | None: - """Return seconds since the bot's most recent comment with ``marker``. - - Fetches comments via `_iter_paginated_json` and delegates the - iteration / author-filter / timestamp logic to - `agent_shin_shared.seconds_since_latest_marker_comment` so the daily - Greptile sweep and the LLM judge use one source of truth for the - "bot already posted X" detection. The wall-clock `now` is resolved - against this module's `dt` so tests that freeze time via - `monkeypatch.setattr(triage_module, "dt", ...)` still apply. - """ - return seconds_since_latest_marker_comment( - _iter_paginated_json(f"repos/{repo}/issues/{number}/comments"), - marker=marker, - bot_login=bot_login, - now=dt.datetime.now(dt.timezone.utc), - ) - - -def seconds_since_last_reconsider_verdict( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since the bot's most recent reconsider verdict comment. - - Detects comments by matching the HTML marker `RECONSIDER_COMMENT_MARKER` - appended by `format_reopen_comment` and - `format_reconsider_still_failing_comment`. Returns None when the bot - has never posted a reconsider verdict on this PR/issue (or when the - only matching comments are missing a `created_at` timestamp, which - shouldn't happen on a real GitHub response). - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=RECONSIDER_COMMENT_MARKER, bot_login=bot_login - ) - - -def seconds_since_last_grace_warning( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since the bot's most recent grace-period warning. - - Detects warning comments by matching the HTML marker - `GRACE_COMMENT_MARKER` appended by `format_grace_warning_pr_comment` - and `format_grace_warning_issue_comment`. Returns None when no - grace warning has ever been posted on this PR/issue — that's the - "first low-quality detection" signal that drives the warning path. - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=GRACE_COMMENT_MARKER, bot_login=bot_login - ) - - -def seconds_since_last_agent_shin_close( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since Agent Shin's most recent auto-close comment. - - Detects close comments by matching `AGENT_SHIN_CLOSE_MARKER` (stamped by - `format_pr_close_comment` / `format_issue_close_comment`). Returns None - when Agent Shin has never closed this PR/issue — the signal - `was_closed_by_agent_shin` uses to keep the reconsider reopen path gated - against closures performed by other workflows sharing the bot identity. - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=AGENT_SHIN_CLOSE_MARKER, bot_login=bot_login - ) - - -# --------------------------------------------------------------------------- -# Author classification - - -def is_internal_contributor(item: dict) -> bool: - """Return True if the PR/issue author should be exempted from triage. - - Fail-safe: if `author_association` is missing or empty (which should never - happen on a successful GitHub REST response but is possible on schema - changes or partial responses), treat the author as INTERNAL so the - destructive close path never fires on an unknown contributor. This matches - the sibling `is_external_pr_author` in `close_low_quality_prs.py`. - """ - login = ((item.get("user") or {}).get("login") or "").lower() - if login.endswith("[bot]") or login in {"dependabot", "github-actions"}: - return True - association = (item.get("author_association") or "").upper() - if not association or association in INTERNAL_ASSOCIATIONS: - return True - return False - - -# --------------------------------------------------------------------------- -# Greptile score + age helpers (`extract_greptile_score`, `parse_iso8601`) -# live in `agent_shin_shared` — they're imported at the top of this module -# so both `triage_with_llm.py` and `close_low_quality_prs.py` share a -# single source of truth for the Confidence-Score regex and ISO-8601 -# parsing. - - -# --------------------------------------------------------------------------- -# Prompt construction - - -def strip_html_comments(text: str) -> str: - """Remove HTML comments — template placeholder text shouldn't fool the judge.""" - return HTML_COMMENT_PATTERN.sub("", text or "") - - -def has_linked_issue(text: str) -> bool: - """Heuristic: does this body link to an open issue (Fixes #123 etc.)?""" - return bool(LINKED_ISSUE_PATTERN.search(strip_html_comments(text or ""))) - - -def build_pr_prompt(*, title: str, body: str) -> str: - cleaned_body = strip_html_comments(body or "").strip() or "(empty)" - # Dedent the static template *before* interpolating dynamic fields so that - # multi-line bodies (whose 2nd+ lines start at column 0) don't defeat the - # common-indent computation in textwrap.dedent. - template = textwrap.dedent(""" - You are "Agent Shin", the OSS triage bot for the LiteLLM open-source - repository (BerriAI/litellm). Decide whether this external pull request - meets the project's contribution standards. - - A PR PASSES triage only if BOTH (1) AND (2) are satisfied. A linked - issue alone is NOT enough — it covers context, not proof. - - (1) CONTEXT — the PR provides AT LEAST ONE of: - (a) A link to a related GitHub issue. Acceptable forms: - "Fixes #1234", "Closes #1234", "Resolves #1234", - "Refs https://github.com/BerriAI/litellm/issues/1234". A - bare "#1234" without a closing keyword counts only if it - is clearly the related issue (not a passing mention). - (b) A clear problem description in the body (what bug or - missing feature this addresses, beyond the title) AND - expected vs. actual behavior (or, for features, "what's - possible now vs. with this PR"). - - (2) END-TO-END QA PROOF: the PR body contains AT LEAST ONE of: - (a) A screen recording / video showing the behavior before - and after the change (the bug reproducing, then the fix - working). For a brand-new feature with no meaningful - "before", a recording of it working end-to-end is fine. - (b) A screenshot (or before/after screenshots) showing the - fix or feature working. - (c) Specific commands that were actually run (curl, python, - a CLI invocation, etc.) PAIRED WITH their real - output, demonstrating the change works end-to-end against - the real system. Commands whose external dependencies - (LLM provider, DB, network) are mocked or stubbed do NOT - satisfy (2c); they are not end-to-end. - - `has_qa_proof` must be set to `true` only when (2a), (2b), - or a non-mocked (2c) is actually present in the body. If the - only "proof" is mocked tests, `has_qa_proof` is `false` and - the verdict is "fail". - - The following do NOT count as QA proof: - - Generic claims like "I tested it", "works locally", "all - tests pass", or a checked "I added tests" checkbox with no - output shown. - - A description of what tests exist or were added, without - their actual output in the PR body. - - `pytest` (or any test runner) executed against the - repository's own unit tests. Those mock the LLM provider, - DB, and network, so they are NOT end-to-end and never - satisfy (2), no matter how much passing output is pasted. - - A linked issue. The linked issue is context (1a), never - proof (2). - - FAIL the PR if EITHER (1) or (2) is missing. Do not bias toward PASS: - if QA proof is absent, the verdict is "fail" even when the rest of - the PR is well-written. - - Respond with a single JSON object, no prose: - - {{ - "verdict": "pass" | "fail", - "linked_issue": boolean, - "has_problem_description": boolean, - "has_expected_vs_actual": boolean, - "has_qa_proof": boolean, - "qa_proof_type": "video" | "screenshot" | "commands_with_output" | "none", - "missing": ["plain-english strings naming what is missing"], - "explanation": "1-2 sentence reasoning for the team to skim" - }} - - --- - PR title: {title} - - PR body: - --- - {cleaned_body} - --- - """).strip() - return template.format(title=title, cleaned_body=cleaned_body) - - -def build_issue_prompt(*, title: str, body: str) -> str: - cleaned_body = strip_html_comments(body or "").strip() or "(empty)" - # Dedent the static template *before* interpolating dynamic fields so that - # multi-line bodies (whose 2nd+ lines start at column 0) don't defeat the - # common-indent computation in textwrap.dedent. - template = textwrap.dedent(""" - You are "Agent Shin", the OSS triage bot for the LiteLLM open-source - repository (BerriAI/litellm). Decide whether this GitHub issue meets - the project's reporting standards. - - For a BUG REPORT the issue PASSES triage only when it contains BOTH: - (1) END-TO-END EVIDENCE OF THE BUG (the "before"; set - `has_repro=true` only when this is present): AT LEAST ONE of: - (a) A screen recording / video of the bug happening. - (b) A screenshot of the bug. - (c) The exact command(s) actually run (curl, python, a CLI - invocation, etc.) PAIRED WITH their real output, traceback, - or logs showing the failure against the real system. - Commands whose external dependencies (LLM provider, DB, - network) are mocked or stubbed do NOT count. - Prose-only "steps to reproduce" with no run output, video, or - screenshot do NOT satisfy (1). An unfilled template scaffold - (bare headings such as "Version or commit:" with nothing under - them, empty numbered lists) counts as absent, not as evidence. - (2) Expected vs. actual behavior (`has_expected_vs_actual`). - - FAIL the bug report if either (1) or (2) is missing. Do not bias - toward PASS: if the bug isn't demonstrated end-to-end, the verdict is - "fail" even when the report is well-written. - - For a FEATURE REQUEST the issue PASSES triage only when it contains - ALL of: - - A clear description of the proposed feature (what should LiteLLM do - that it does not today). - - Motivation / use case with a concrete example (config, API call, - UI flow, or scenario showing what's blocked today). - - END-TO-END EVIDENCE OF THE DEAD-END (set - `has_dead_end_evidence=true` only when this is present): a video, - a screenshot, or the exact command(s) actually run paired with - their real output, showing the point where the flow stops today. - Mocked or stubbed dependencies do NOT count, and an unfilled - template scaffold (bare headings, empty numbered lists) counts as - absent. - - For an issue that is neither a bug report nor a feature request (a - question, support request, or discussion), PASS as long as it has a - clear, specific ask and is not empty or template placeholder text. - - Respond with a single JSON object, no prose: - - {{ - "verdict": "pass" | "fail", - "kind": "bug" | "feature" | "other", - "has_repro": boolean, - "has_expected_vs_actual": boolean, - "has_motivation_example": boolean, - "has_dead_end_evidence": boolean, - "missing": ["plain-english strings naming what is missing"], - "explanation": "1-2 sentence reasoning for the team to skim" - }} - - --- - Issue title: {title} - - Issue body: - --- - {cleaned_body} - --- - """).strip() - return template.format(title=title, cleaned_body=cleaned_body) - - -# --------------------------------------------------------------------------- -# LLM call + verdict parsing - - -def call_llm_judge( - prompt: str, *, model: str, api_key: str, base_url: str | None -) -> str: - """Call an OpenAI-compatible chat completions endpoint. Returns raw text.""" - # Import inside the function so unit tests that monkey-patch this never - # need the openai package installed. - from openai import OpenAI - - client = ( - OpenAI(api_key=api_key, base_url=base_url) - if base_url - else OpenAI(api_key=api_key) - ) - kwargs: dict[str, Any] = { - "model": model, - "messages": [{"role": "user", "content": prompt}], - "temperature": 0, - "response_format": {"type": "json_object"}, - } - # gpt-5.x reasoning models reject `temperature != 1` unless - # `reasoning_effort` is explicitly "none". Set it via `extra_body` so this - # works across openai SDK versions regardless of whether the SDK natively - # types `reasoning_effort` as a top-level chat-completions param yet. - if model.lower().startswith(GPT5_FAMILY_PREFIX): - kwargs["extra_body"] = {"reasoning_effort": "none"} - response = client.chat.completions.create(**kwargs) - return response.choices[0].message.content or "" - - -def parse_verdict(raw: str) -> dict: - """Parse the LLM's JSON response. Tolerates ```json fences and stray text.""" - if not raw: - raise ValueError("empty LLM response") - text = raw.strip() - if text.startswith("```"): - text = re.sub(r"^```(?:json)?\s*", "", text) - text = re.sub(r"\s*```$", "", text) - try: - return json.loads(text) - except json.JSONDecodeError: - match = re.search(r"\{.*\}", text, re.DOTALL) - if not match: - raise ValueError(f"could not extract JSON from LLM response: {raw[:200]}") - return json.loads(match.group(0)) - - -# --------------------------------------------------------------------------- -# Comment composition - - -def _format_missing(missing: list[str]) -> str: - if not missing: - return "- (see explanation below)" - return "\n".join(f"- {m}" for m in missing) - - -# Rubric items the judge can mark present. The first element of each tuple is -# the verdict-JSON boolean field, the second is the human-readable label we -# render in the "what you got right" section of close / grace-warning comments. -_PR_PRESENT_LABELS: tuple[tuple[str, str], ...] = ( - ("linked_issue", "Linked a related GitHub issue"), - ("has_problem_description", "Clear problem description"), - ("has_expected_vs_actual", "Expected vs. actual behavior"), - ("has_qa_proof", "End-to-end QA proof"), -) - -# Issue rubric labels grouped by `kind`. The judge sets `kind` to one of -# {"bug", "feature", "other"}; when "other" we render both groups so we don't -# silently drop a present-flag the judge actually set to True. -_ISSUE_BUG_LABELS: tuple[tuple[str, str], ...] = ( - ( - "has_repro", - "End-to-end evidence of the bug (video, screenshot, or command + real output)", - ), - ("has_expected_vs_actual", "Expected vs. actual behavior"), -) -_ISSUE_FEATURE_LABELS: tuple[tuple[str, str], ...] = ( - ("has_motivation_example", "Motivation and concrete example"), - ( - "has_dead_end_evidence", - "End-to-end evidence of the dead-end (video, screenshot, or command + real output)", - ), -) - - -def _format_present_for_pr(verdict: dict) -> list[str]: - """Human-readable rubric items the judge confirmed are present on a PR. - - Drives the "what you got right" section in close / grace-warning comments. - The user gave explicit feedback: contributors should see what they nailed - *before* the list of gaps, so the comment doesn't read as pure rejection. - """ - return [label for field, label in _PR_PRESENT_LABELS if verdict.get(field)] - - -def _format_present_for_issue(verdict: dict) -> list[str]: - """Human-readable rubric items the judge confirmed are present on an issue. - - Branches on the judge's `kind` field. For `"other"` (or missing kind) we - render the union so a present-flag isn't dropped just because the judge - couldn't classify the issue cleanly. - """ - kind = (verdict.get("kind") or "").lower() - groups: list[tuple[tuple[str, str], ...]] = [] - if kind in ("bug", "other", ""): - groups.append(_ISSUE_BUG_LABELS) - if kind in ("feature", "other", ""): - groups.append(_ISSUE_FEATURE_LABELS) - out: list[str] = [] - for group in groups: - for field, label in group: - if verdict.get(field) and label not in out: - out.append(label) - return out - - -def _format_present_block(items: list[str]) -> str: - """Render the optional "what you got right" block. Empty string when the - judge didn't confirm anything as present — better to omit the section - entirely than to show "What you got right: (nothing)". - """ - if not items: - return "" - bullets = "\n".join(f"- ✅ {item}" for item in items) - return f"**What you got right:**\n\n{bullets}\n\n" - - -def format_pr_close_comment(verdict: dict) -> str: - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_pr(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the description against our " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md). " - "Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "**Closing this PR isn't a rejection of the change.** We want the open-PR list to " - "mirror what a maintainer can act on *right now*, so contributors don't get lost in a " - 'backlog. A closed PR is a soft "park this for later"; your work is still here, ' - "the diff is still here, and getting it reopened is one comment away. Take your time.\n" - "\n" - "**To bring this PR back:**\n" - "\n" - "- Update the description with the missing pieces, then comment `@agent-shin reconsider` " - "on this PR. I'll re-evaluate and reopen if it now passes.\n" - "- Or **Open a new PR** with the same fix and the updated description. GitHub doesn't " - "always let external contributors reopen a bot-closed PR, so a fresh PR is the most " - "reliable path back into the review queue.\n" - "- If Greptile's most recent score on this PR was below 4/5, comment `@greptileai` to " - "request a fresh review; that **still works even after the PR is closed**, and a " - "stronger score is one of the signals that lifts the PR back into the queue. A low " - "Greptile score isn't a blocker.\n" - "\n" - '**What "end-to-end QA proof" means**, since it\'s the most common gap: at least one ' - "of a short before/after screen recording / video (the bug reproducing, then the fix " - "working; for a brand-new feature, a recording of it working end-to-end), a screenshot " - "(or before/after screenshots) of it working, or the exact commands you ran paired " - "with their **real output** against the real system. Running `pytest` on the repo's " - "unit tests doesn't count; those mock the LLM provider, DB, and network, so they " - "aren't end-to-end. Output from a real, no-mocks integration run is what we look " - "for. A linked issue alone isn't enough either: it covers context, not proof. See " - "[the full rubric](https://docs.litellm.ai/blog/agent-shin-triage#the-rubric-for-pull-requests).\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, comment " - "`@agent-shin reconsider` or ping a maintainer; they'll override me.)_" - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def format_issue_close_comment(verdict: dict) -> str: - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_issue(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for filing this! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the issue against our reporting checklist. Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "**Closing this isn't us saying the bug isn't real or the request isn't useful.** We " - "want the open-issue list to mirror what a maintainer can act on *right now*, so " - "reports like yours don't get buried in a backlog. A closed issue is a soft \"park " - 'this for later"; your report is still here, and getting it reopened is one comment ' - "away. Take your time.\n" - "\n" - "**To bring this issue back:**\n" - "\n" - "1. Edit the issue description to add the missing pieces:\n" - " - For **bug reports**: end-to-end evidence of the bug (a screen recording / " - "video, a screenshot, or the exact commands you ran with their real output / " - "traceback) plus expected vs. actual behavior. Written steps with no run output, " - "video, or screenshot don't count, and mocked or stubbed runs don't count.\n" - " - For **feature requests**: a concrete description of what should change, a " - "use case and example (config / API call / UI flow), plus end-to-end evidence of " - "the dead-end (a video, a screenshot, or the exact commands you ran with their " - "real output showing where the flow stops today). Mocked or stubbed runs don't " - "count.\n" - "2. Comment `@agent-shin reconsider`. I'll re-run triage and reopen the issue if it " - "now meets the bar. (GitHub doesn't let external authors reopen an issue a maintainer " - "or bot closed, so the comment-based reconsider is the reliable path.)\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, comment " - "`@agent-shin reconsider` or ping a maintainer; they'll override me.)_" - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def format_grace_warning_pr_comment(verdict: dict) -> str: - """Comment posted on the FIRST low-quality detection — gives the - contributor a 2-hour grace window to fix the PR before the next - triage run actually closes it. - - This is the "before-close" warning. On the second triage run, if the - grace marker is older than `GRACE_PERIOD_SECONDS` AND the PR still - fails the rubric, the close path runs (which posts - `format_pr_close_comment` and closes the PR). - """ - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_pr(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the description against our " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md). " - "Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "If the description isn't updated in the next **2 hours**, I'll auto-close this PR. " - "That's **not** us saying we don't care about the change; we want the open-PR list to " - "mirror what a maintainer can act on *right now*, so contributors don't get lost in a " - 'backlog. A closed PR is a soft "park this for later," not a rejection. Take your ' - "time; everything below still works after the close.\n" - "\n" - "**During the grace period:** just update the PR description with the missing pieces. " - "No need to ping me; I'll re-check on the next sweep and skip the auto-close if it " - "now passes. See " - "[what counts as QA proof](https://docs.litellm.ai/blog/agent-shin-triage#the-rubric-for-pull-requests) " - "for the full rubric (a linked issue alone isn't enough; it covers context, not proof).\n" - "\n" - "**If the PR does get auto-closed in 2 hours, you still have easy recovery paths:**\n" - "\n" - "- Comment `@agent-shin reconsider` after updating the description. I'll re-evaluate " - "and reopen the PR if it now passes.\n" - "- Comment `@greptileai` to request a fresh Greptile review; that **still works even " - "after the PR is closed**, and a stronger score is one of the signals that lifts the " - "PR back into the queue. So a low Greptile score isn't a blocker either.\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, ping a " - "maintainer; they'll override me.)_\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -def format_grace_warning_issue_comment(verdict: dict) -> str: - """Issue analogue of `format_grace_warning_pr_comment`.""" - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_issue(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for filing this! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the issue against our reporting checklist. Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "If the issue isn't updated in the next **2 hours**, I'll auto-close it. That's **not** us " - "saying the bug isn't real or the request isn't useful; we want the open-issue list " - "to mirror what a maintainer can act on *right now*, so reports like yours don't get " - 'buried in a backlog. A closed issue is a soft "park this for later," not a ' - "rejection. Take your time; reopening is one comment away.\n" - "\n" - "**During the grace period:** just edit the issue description with the missing " - "pieces. No need to ping me; I'll re-check on the next sweep and skip the auto-close " - "if it now passes.\n" - "\n" - "Missing pieces, depending on what this is:\n" - "\n" - "- For **bug reports**: end-to-end evidence of the bug (a screen recording / video, a " - "screenshot, or the exact commands you ran with their real output / traceback) plus " - "expected vs. actual behavior. Written steps with no run output don't count, and " - "mocked or stubbed runs don't count.\n" - "- For **feature requests**: a concrete description of what should change, a use " - "case and example (config / API call / UI flow), plus end-to-end evidence of the " - "dead-end (a video, a screenshot, or the exact commands you ran with their real " - "output showing where the flow stops today). Mocked or stubbed runs don't count.\n" - "\n" - "**If the issue does get auto-closed in 2 hours**, comment `@agent-shin reconsider` " - "and I'll re-evaluate. If it now meets the bar, I'll reopen the issue.\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, ping a " - "maintainer; they'll override me.)_\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -# --------------------------------------------------------------------------- -# Step-summary helpers - - -def write_step_summary(content: str) -> None: - """When running inside GitHub Actions, append to the step summary file.""" - path = os.environ.get("GITHUB_STEP_SUMMARY") - if not path: - return - try: - with open(path, "a", encoding="utf-8") as handle: - handle.write(content) - if not content.endswith("\n"): - handle.write("\n") - except OSError as exc: - print(f"warn: failed to write step summary: {exc}", file=sys.stderr) - - -# --------------------------------------------------------------------------- -# Core orchestration - - -def format_reopen_comment(kind: str) -> str: - """Comment posted when Agent Shin reopens after a successful reconsider.""" - noun = "PR" if kind == "pr" else "issue" - # The trailing HTML marker is used by `seconds_since_last_reconsider_verdict` - # to enforce a cooldown between repeated `@agent-shin reconsider` triggers. - # Keep the marker on its own line so it doesn't disturb the rendered text. - return ( - f"♻️ **Re-evaluated and reopened.** Thanks for updating the {noun}!\n" - "\n" - "Agent Shin re-ran triage on the latest description and it now meets " - "the bar. A maintainer will take another look soon; please don't " - f"close this {noun} again unless asked to.\n" - "\n" - "_(If a maintainer ends up closing this for non-rubric reasons, that " - "decision stands; comment `@agent-shin reconsider` again only if you " - "have substantively new information.)_\n" - "\n" - f"{RECONSIDER_COMMENT_MARKER}" - ) - - -def format_reconsider_still_failing_comment(kind: str, verdict: dict) -> str: - """Comment posted when reconsider re-runs triage but the verdict is still fail.""" - missing_lines = _format_missing(verdict.get("missing") or []) - explanation = verdict.get("explanation") or "" - noun = "PR" if kind == "pr" else "issue" - # The trailing HTML marker is used by `seconds_since_last_reconsider_verdict` - # to enforce a cooldown between repeated `@agent-shin reconsider` triggers. - return ( - f"⏸️ **Re-evaluated; this {noun} still doesn't meet the rubric.**\n" - "\n" - "Agent Shin re-ran triage on the current description but is still " - "missing:\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "Update the description with the missing pieces and comment " - "`@agent-shin reconsider` again, or ping a maintainer if you think " - "I got this wrong.\n" - "\n" - "_(I'm an LLM and I'm not infallible.)_\n" - "\n" - f"{RECONSIDER_COMMENT_MARKER}" - ) - - -# --------------------------------------------------------------------------- -# Review gate — "ready for review" label lifecycle - -_UNSET = object() - - -def _combine_missing( - verdict: dict, greptile_score: int | None, min_score: int -) -> list[str]: - """Merge the LLM rubric's `missing` list with a Greptile-score shortfall.""" - missing = list(verdict.get("missing") or []) - if greptile_score is not None and greptile_score < min_score: - missing.insert( - 0, - f"Greptile's most recent review scored this PR {greptile_score}/5 " - f"(below the {min_score}/5 bar)", - ) - return missing or ["(see explanation below)"] - - -def _has_marker( - comments: Iterable[dict], marker: str, *, bot_login: str | None = None -) -> bool: - """Return True iff the bot itself posted a comment containing ``marker``. - - Filters by author so a contributor who quotes the marker (e.g. via - GitHub's "Quote reply" feature, which preserves HTML comments in - raw markdown) is not mistaken for a bot action — that would - silently suppress notifications or change which "recovered" wording - is selected. Matches the author-filter pattern used by the sibling - `_seconds_since_latest_marker_comment` helper. - """ - expected_login = ( - bot_login - or os.environ.get("AGENT_SHIN_BOT_LOGIN") - or AGENT_SHIN_DEFAULT_BOT_LOGIN - ).lower() - for comment in comments: - author = ((comment.get("user") or {}).get("login") or "").lower() - if author != expected_login: - continue - if marker in (comment.get("body") or ""): - return True - return False - - -def format_ready_for_review_comment( - verdict: dict, - greptile_score: int | None, - min_greptile_score: int = DEFAULT_MIN_GREPTILE_SCORE, -) -> str: - """Posted the first time a PR clears the bar (label added).""" - score_line = ( - f" Greptile scored it **{greptile_score}/5**." - if greptile_score is not None - else "" - ) - explanation = verdict.get("explanation") or "" - return ( - "✅ **Triage passed, tagging `ready for review`.**\n" - "\n" - "Agent Shin checked this PR against the " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md) " - "and it clears the bar (a linked issue, or a clear problem description " - f"+ expected vs. actual + QA proof).{score_line}\n" - "\n" - f"> {explanation}\n" - "\n" - "A maintainer will take it from here. If a later re-check finds the PR " - f"has regressed (Greptile drops below {min_greptile_score}/5, " - "the QA proof is removed, etc.) I'll pull the tag and comment with " - "what's missing; fix it and the tag comes back automatically.\n" - f"{READY_MARKER}" - ) - - -def format_all_clear_comment(verdict: dict, greptile_score: int | None) -> str: - """Posted when a PR recovers after a regression (label re-added).""" - score_line = ( - f" Greptile is back to **{greptile_score}/5**." - if greptile_score is not None - else "" - ) - explanation = verdict.get("explanation") or "" - return ( - "✅ **All clear again, re-adding `ready for review`.**\n" - "\n" - "Thanks for addressing the earlier feedback. On re-check this PR meets " - f"the contribution bar once more.{score_line}\n" - "\n" - f"> {explanation}\n" - "\n" - "A maintainer will take another look.\n" - f"{READY_MARKER}" - ) - - -def format_regression_comment( - missing: list[str], explanation: str, grace_days: int -) -> str: - """Posted when a previously-tagged PR regresses (label removed, PR stays open). - - Discloses the same ``grace_days`` deadline the state machine enforces: - once that window elapses with the PR still failing, the close path fires. - Hiding the deadline behind a bare "stays open" would surprise contributors - with an auto-close they were never warned about. - """ - window = "24 hours" if grace_days == 1 else f"{grace_days} days" - return ( - "⚠️ **Removing the `ready for review` tag.**\n" - "\n" - "On a re-check this PR no longer meets the contribution bar. What's " - "missing now:\n" - "\n" - f"{_format_missing(missing)}\n" - "\n" - f"> {explanation}\n" - "\n" - f"The PR stays open for ~{window}; address the points above and Agent " - 'Shin will post an "all clear" comment and re-add the tag ' - "automatically. If the points still aren't addressed after that " - "window, the PR is auto-closed; that's not a rejection, and you can " - "comment `@agent-shin reconsider` to have it re-evaluated and reopened " - "once it passes.\n" - f"{REGRESSED_MARKER}" - ) - - -def format_within_grace_comment( - missing: list[str], explanation: str, grace_days: int -) -> str: - """Posted once while a failing PR is still inside its grace window.""" - window = "24 hours" if grace_days == 1 else f"{grace_days} days" - return ( - "🚅 Hi, thanks for the PR! This is **Agent Shin**, the automated triage " - "bot. This PR doesn't quite meet the contribution bar yet:\n" - "\n" - f"{_format_missing(missing)}\n" - "\n" - f"> {explanation}\n" - "\n" - f"You have ~{window} from when this PR was opened to add the missing " - "pieces; just update the description and I'll re-check on the next " - "sweep. Once it passes I'll tag it `ready for review`. If it does get " - "auto-closed, that's not a rejection; comment `@agent-shin reconsider` " - "and I'll re-evaluate and reopen if it now passes.\n" - f"{WITHIN_GRACE_MARKER}" - ) - - -def review_gate( - *, - repo: str, - number: int, - close: bool, - model: str, - judge: Any = None, - greptile_score: Any = _UNSET, - comments: Any = _UNSET, - now: dt.datetime | None = None, - grace_days: int = DEFAULT_GRACE_DAYS, - min_greptile_score: int = DEFAULT_MIN_GREPTILE_SCORE, - label: str = READY_FOR_REVIEW_LABEL, - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> dict: - """Reconcile the `ready for review` label with a PR's current quality. - - A PR is *passing* when it clears BOTH gates: the LLM rubric (linked issue, - or problem description + expected/actual + QA proof) AND Greptile's most - recent confidence score (>= ``min_greptile_score``; absence of a score is - not held against the PR). The gate then drives a small state machine, using - the label itself as the persisted state so comments fire only on - transitions (never on every scheduled run): - - passing, untagged -> add label + "ready for review" / "all clear" - passing, tagged -> noop-passing - not passing, tagged -> remove label + regression comment (stays open) - not passing, untagged, old -> close + comment (past the grace window) - not passing, untagged, new -> one-time "what's missing" notice (within grace) - - ``close`` gates every destructive side effect: with ``close=False`` the - function returns a ``would-*`` preview and touches nothing, mirroring the - dry-run contract of :func:`triage`. ``judge``/``greptile_score``/ - ``comments``/``now`` are injectable for tests; in production they are - resolved from the OpenAI judge, the PR's Greptile comment, the live comment - list, and the wall clock respectively. - """ - item = fetch_pr(repo, number) - - title = item.get("title") or "" - body = item.get("body") or "" - login = (item.get("user") or {}).get("login") or "" - association = item.get("author_association") or "" - state = item.get("state") or "" - # GitHub label names are case-insensitive; compare lowercased so a repo - # that already has e.g. "Ready for Review" is recognized as the same - # label as our READY_FOR_REVIEW_LABEL constant ("ready for review"). - labels_now = {(lbl.get("name") or "").lower() for lbl in (item.get("labels") or [])} - label_key = label.lower() - created_raw = item.get("created_at") or "" - - base_result = { - "kind": "pr", - "number": number, - "title": title, - "author": login, - "author_association": association, - "state": state, - "labeled": label_key in labels_now, - "review_gate": True, - } - - if state != "open": - return {**base_result, "action": "skip-not-open"} - - if allowlist: - if login.lower() not in allowlist: - return {**base_result, "action": "skip-not-allowlisted"} - elif is_internal_contributor(item): - return {**base_result, "action": "skip-internal-author"} - - # Resolve the comment list once — used for both the Greptile score and the - # marker-based dedup below. - if comments is _UNSET: - comments = list(_iter_paginated_json(f"repos/{repo}/issues/{number}/comments")) - - # --- rubric verdict: linked-issue short-circuit, else the LLM judge ------- - if has_linked_issue(body): - verdict = { - "verdict": "pass", - "linked_issue": True, - "missing": [], - "explanation": "Linked-issue regex matched; LLM was not called.", - } - rubric_pass = True - else: - prompt = build_pr_prompt(title=title, body=body) - if judge is None: - api_key = os.environ.get("OPENAI_API_KEY") - if not api_key: - return {**base_result, "action": "skip-no-llm-key"} - base_url = os.environ.get("OPENAI_BASE_URL") or None - - def judge(p: str) -> str: - return call_llm_judge( - p, model=model, api_key=api_key, base_url=base_url - ) - - try: - verdict = parse_verdict(judge(prompt)) - except Exception as exc: # noqa: BLE001 - judge errors must never act - return {**base_result, "action": "skip-llm-error", "error": str(exc)} - rubric_pass = (verdict.get("verdict") or "").lower() == "pass" - - # --- Greptile score ------------------------------------------------------- - if greptile_score is _UNSET: - extraction = extract_greptile_score(comments) - greptile_score = extraction[0] if extraction else None - greptile_ok = greptile_score is None or greptile_score >= min_greptile_score - passing = rubric_pass and greptile_ok - - # --- age ------------------------------------------------------------------ - age_days = None - if created_raw: - reference = now or dt.datetime.now(dt.timezone.utc) - age_days = (reference - parse_iso8601(created_raw)).days - - label_present = label_key in labels_now - explanation = verdict.get("explanation") or "" - # When the rubric short-circuited to pass (linked-issue regex) but - # Greptile dragged the PR below the bar, the synthetic verdict's - # explanation ("LLM was not called") would mislead a contributor reading - # the regression / close comment. Surface the real reason instead. - if rubric_pass and not greptile_ok: - explanation = ( - f"Greptile's most recent review scored this PR " - f"{greptile_score}/5 (below the {min_greptile_score}/5 bar)." - ) - verdict = {**verdict, "explanation": explanation} - base_result = { - **base_result, - "verdict": verdict, - "greptile_score": greptile_score, - "passing": passing, - "age_days": age_days, - } - - if passing: - if label_present: - return {**base_result, "action": "noop-passing"} - recovered = _has_marker(comments, REGRESSED_MARKER) - comment = ( - format_all_clear_comment(verdict, greptile_score) - if recovered - else format_ready_for_review_comment( - verdict, greptile_score, min_greptile_score - ) - ) - if not close: - return {**base_result, "action": "would-label-ready", "comment": comment} - post_comment(repo, number, comment) - add_label(repo, number, label) - return {**base_result, "action": "labeled-ready", "comment": comment} - - missing = _combine_missing(verdict, greptile_score, min_greptile_score) - - if label_present: - comment = format_regression_comment(missing, explanation, grace_days) - if not close: - return {**base_result, "action": "would-remove-label", "comment": comment} - remove_label(repo, number, label) - post_comment(repo, number, comment) - return {**base_result, "action": "label-removed-regressed", "comment": comment} - - # Not passing and not tagged. If the PR was previously tagged and then - # regressed (we removed the label and posted REGRESSED_MARKER), honor the - # "PR stays open — fix it and the tag comes back" promise from - # `format_regression_comment` and skip the close path. Without this guard, - # any PR older than `grace_days` would be closed on the next evaluation, - # giving the contributor no realistic window to address the regression. - # - # The promise has a deliberate expiration: once `grace_days` have elapsed - # since the regression notice, fall through to the close path so a PR that - # was abandoned post-regression doesn't sit open forever. - if _has_marker(comments, REGRESSED_MARKER): - reference = now or dt.datetime.now(dt.timezone.utc) - seconds_since_regression = seconds_since_latest_marker_comment( - comments, marker=REGRESSED_MARKER, now=reference - ) - grace_seconds = grace_days * 86400 - if seconds_since_regression is None or seconds_since_regression < grace_seconds: - return {**base_result, "action": "regressed-already-notified"} - - # Not passing and not tagged: close if past the grace window, else notify once. - if age_days is not None and age_days >= grace_days: - comment = format_pr_close_comment({**verdict, "missing": missing}) - if not close: - return {**base_result, "action": "would-close", "comment": comment} - post_comment(repo, number, comment) - close_pr(repo, number) - return {**base_result, "action": "closed", "comment": comment} - - if _has_marker(comments, WITHIN_GRACE_MARKER): - return {**base_result, "action": "within-grace-already-notified"} - comment = format_within_grace_comment(missing, explanation, grace_days) - if not close: - return { - **base_result, - "action": "would-notify-within-grace", - "comment": comment, - } - post_comment(repo, number, comment) - return {**base_result, "action": "within-grace-notified", "comment": comment} - - -def triage( - *, - repo: str, - kind: str, - number: int, - close: bool, - model: str, - judge: Any = None, - print_prompt: bool = False, - reconsider: bool = False, - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> dict: - """Triage a single PR or issue. Returns a result dict for logging/tests. - - `judge` is an optional callable `(prompt) -> str` for tests / dry-run with - a stub. In production, leave it None and the script uses `call_llm_judge`. - - When `reconsider=True`, the closed-state guard is skipped and a - fail-but-no-comment is replaced with a "still failing" comment + leave - closed; a pass triggers `reopen_pr`/`reopen_issue` plus a reopen comment. - Reconsider mode is intended for the `@agent-shin reconsider` comment - trigger. Like regular triage, `close=False` keeps reconsider in dry-run - (returns `would-reopen` / `would-reconsider-still-failing` so a local - operator can preview without write side effects); the workflow only - passes `--close` when `AGENT_SHIN_ENABLED=true`. - - Reconsider mode adds two extra safety guards on top of the regular - triage skip-internal-author check: - - 1. **Bot-closed guard.** Only reopens if the most recent close was - performed by the bot identity (default `github-actions[bot]`). - This stops a contributor from using `@agent-shin reconsider` to - override a maintainer's close for non-rubric reasons. - 2. **Rate-limit guard.** If the bot has already posted a reconsider - verdict on this PR/issue within `RECONSIDER_RATE_LIMIT_SECONDS`, - skip — repeated triggers from the same contributor shouldn't burn - CI minutes or LLM budget. - """ - fetcher = {"pr": fetch_pr, "issue": fetch_issue}[kind] - item = fetcher(repo, number) - - title = item.get("title") or "" - body = item.get("body") or "" - login = (item.get("user") or {}).get("login") or "" - association = item.get("author_association") or "" - state = item.get("state") or "" - - base_result = { - "kind": kind, - "number": number, - "title": title, - "author": login, - "author_association": association, - "state": state, - "reconsider": reconsider, - } - - # Reconsider only makes sense on a closed PR/issue. A "reconsider on an - # open PR" is a no-op (the regular triage flow already evaluates open - # PRs); return a clear skip so the workflow can short-circuit. - if reconsider: - if state != "closed": - return {**base_result, "action": "skip-not-closed"} - else: - if state != "open": - return {**base_result, "action": "skip-not-open"} - - if allowlist: - if login.lower() not in allowlist: - return {**base_result, "action": "skip-not-allowlisted"} - elif is_internal_contributor(item): - return {**base_result, "action": "skip-internal-author"} - - # Reconsider-only guards — these run BEFORE the LLM call so a - # maintainer-closed PR / rate-limited trigger never spends LLM budget. - if reconsider: - if not was_closed_by_agent_shin(repo, number): - return {**base_result, "action": "skip-not-bot-closed"} - age = seconds_since_last_reconsider_verdict(repo, number) - if age is not None and age < RECONSIDER_RATE_LIMIT_SECONDS: - return { - **base_result, - "action": "skip-rate-limited", - "rate_limit_age_seconds": age, - "rate_limit_window_seconds": RECONSIDER_RATE_LIMIT_SECONDS, - } - - if kind == "pr": - # Short-circuit: if body very clearly links a related issue, just pass. - if has_linked_issue(body): - base = { - **base_result, - "action": "pass-linked-issue", - "verdict": { - "verdict": "pass", - "linked_issue": True, - "explanation": "Linked-issue regex matched; LLM was not called.", - }, - } - if reconsider: - # Pass-on-reconsider -> reopen the PR with a friendly comment. - reopen_body = format_reopen_comment(kind) - if not close: - return { - **base, - "action": "would-reopen", - "comment": reopen_body, - } - post_comment(repo, number, reopen_body) - reopen_pr(repo, number) - return { - **base, - "action": "reopened", - "comment": reopen_body, - } - return base - prompt = build_pr_prompt(title=title, body=body) - else: - prompt = build_issue_prompt(title=title, body=body) - - if print_prompt: - return {**base_result, "action": "print-prompt", "prompt": prompt} - - if judge is None: - api_key = os.environ.get("OPENAI_API_KEY") - if not api_key: - # No key configured — never take a destructive action. Report skip. - return { - **base_result, - "action": "skip-no-llm-key", - "prompt_preview": prompt[:200], - } - base_url = os.environ.get("OPENAI_BASE_URL") or None - - def judge(p: str) -> str: - return call_llm_judge(p, model=model, api_key=api_key, base_url=base_url) - - try: - raw = judge(prompt) - verdict = parse_verdict(raw) - except Exception as exc: # noqa: BLE001 - judge errors must never close PRs - return {**base_result, "action": "skip-llm-error", "error": str(exc)} - - decision = (verdict.get("verdict") or "").lower() - - if reconsider: - # Reconsider: an explicit `pass` -> reopen + post reopen comment; - # anything else (fail, missing/malformed verdict, typo) -> leave - # closed + post a "still failing" comment so the contributor can - # iterate again. Reopen is destructive, so a flaky/empty verdict - # must not satisfy the gate. - # In dry-run (`close=False`) we return `would-*` actions instead - # of touching GitHub state, mirroring the regular triage flow's - # `would-close`. This lets a local operator preview the outcome - # of `python triage_with_llm.py --reconsider --pr N` without - # risking accidental comments or reopens. - if decision == "pass": - reopen_body = format_reopen_comment(kind) - if not close: - return { - **base_result, - "action": "would-reopen", - "verdict": verdict, - "comment": reopen_body, - } - post_comment(repo, number, reopen_body) - if kind == "pr": - reopen_pr(repo, number) - else: - reopen_issue(repo, number) - return { - **base_result, - "action": "reopened", - "verdict": verdict, - "comment": reopen_body, - } - still_failing = format_reconsider_still_failing_comment(kind, verdict) - if not close: - return { - **base_result, - "action": "would-reconsider-still-failing", - "verdict": verdict, - "comment": still_failing, - } - post_comment(repo, number, still_failing) - return { - **base_result, - "action": "reconsider-still-failing", - "verdict": verdict, - "comment": still_failing, - } - - if decision != "fail": - return {**base_result, "action": "pass-llm", "verdict": verdict} - - # Grace-period flow: on the first low-quality detection, post a warning - # comment instead of closing immediately. On a subsequent triage run - # (manual re-trigger, or the daily `close_low_quality_prs.py` cron - # finding the same PR in its own pass), if `GRACE_PERIOD_SECONDS` has - # elapsed since the warning AND the PR still fails the rubric, close. - grace_age = seconds_since_last_grace_warning(repo, number) - if grace_age is None: - warning_body = ( - format_grace_warning_pr_comment(verdict) - if kind == "pr" - else format_grace_warning_issue_comment(verdict) - ) - if not close: - return { - **base_result, - "action": "would-warn-grace", - "verdict": verdict, - "comment": warning_body, - } - post_comment(repo, number, warning_body) - return { - **base_result, - "action": "warned-grace", - "verdict": verdict, - "comment": warning_body, - } - if grace_age < GRACE_PERIOD_SECONDS: - return { - **base_result, - "action": "skip-in-grace-period", - "verdict": verdict, - "grace_age_seconds": grace_age, - "grace_period_seconds": GRACE_PERIOD_SECONDS, - } - - # The grace window has elapsed. `--close` still gates the destructive - # write so a dry-run preview never posts or closes — the workflow only - # passes `--close` when `AGENT_SHIN_ENABLED=true`, which keeps the bot - # inert by default. - if not close: - return {**base_result, "action": "would-close", "verdict": verdict} - - comment_body = ( - format_pr_close_comment(verdict) - if kind == "pr" - else format_issue_close_comment(verdict) - ) - post_comment(repo, number, comment_body) - if kind == "pr": - close_pr(repo, number) - else: - close_issue(repo, number) - - return { - **base_result, - "action": "closed", - "verdict": verdict, - "comment": comment_body, - } - - -# --------------------------------------------------------------------------- -# CLI - - -def render_summary(result: dict) -> str: - """Render a human-readable summary block (used for stdout + step summary).""" - lines = ["## Agent Shin verdict", ""] - lines.append( - f"- **{result['kind'].upper()} #{result['number']}**: {result.get('title', '')}" - ) - lines.append( - f"- **Author**: `{result.get('author', '')}` ({result.get('author_association', '')})" - ) - lines.append(f"- **State**: {result.get('state', '')}") - lines.append(f"- **Action**: `{result['action']}`") - verdict = result.get("verdict") - if verdict: - lines.append("") - lines.append("```json") - lines.append(json.dumps(verdict, indent=2)) - lines.append("```") - error = result.get("error") - if error: - lines.append("") - lines.append(f"_LLM error: {error}_") - comment = result.get("comment") - if comment: - lines.append("") - lines.append("### Posted comment:") - lines.append("") - lines.append("> " + comment.replace("\n", "\n> ")) - return "\n".join(lines) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--repo", required=True, help="Repository (owner/repo).") - target = parser.add_mutually_exclusive_group(required=True) - target.add_argument("--pr", type=int, help="Pull request number to triage.") - target.add_argument("--issue", type=int, help="Issue number to triage.") - parser.add_argument( - "--close", - action="store_true", - help="Actually post comment + close on fail (default: dry run).", - ) - parser.add_argument( - "--model", - # `os.environ.get("TRIAGE_MODEL", DEFAULT_MODEL)` would return "" when - # GitHub Actions exposes an unset repo variable as an empty-string env - # var, silently bypassing DEFAULT_MODEL and causing every call to fail - # as `skip-llm-error`. The `or` guard collapses empty -> default. - default=os.environ.get("TRIAGE_MODEL") or DEFAULT_MODEL, - help=f"OpenAI-compatible model name (default: {DEFAULT_MODEL}).", - ) - parser.add_argument( - "--print-prompt", - action="store_true", - help="Print the prompt that would be sent to the judge and exit.", - ) - parser.add_argument( - "--reconsider", - action="store_true", - help=( - "Re-run triage on a CLOSED PR/issue and reopen it on pass. " - "Used by the `@agent-shin reconsider` comment-trigger workflow. " - "Only invoke this from a workflow that has already gated on " - "AGENT_SHIN_ENABLED=true and verified the commenter is the " - "PR/issue author or an internal collaborator." - ), - ) - parser.add_argument( - "--review-gate", - action="store_true", - help=( - "Reconcile the `ready for review` label for an OPEN PR: tag on " - "pass, remove the tag + comment on regression, close after the " - "grace window if it never passed. PR-only." - ), - ) - parser.add_argument( - "--grace-days", - type=int, - default=DEFAULT_GRACE_DAYS, - help=( - "Review-gate only: hours/24 a failing, un-tagged PR may stay open " - f"before auto-close (default: {DEFAULT_GRACE_DAYS} = 24h)." - ), - ) - parser.add_argument( - "--min-greptile-score", - type=int, - default=DEFAULT_MIN_GREPTILE_SCORE, - choices=range(1, 6), - help=( - "Review-gate only: Greptile score below which a PR counts as not " - f"passing (default: {DEFAULT_MIN_GREPTILE_SCORE} -> <4/5 regresses)." - ), - ) - args = parser.parse_args() - - kind = "pr" if args.pr is not None else "issue" - number = args.pr if args.pr is not None else args.issue - - if args.review_gate: - if kind != "pr": - parser.error("--review-gate applies to pull requests only (use --pr).") - result = review_gate( - repo=args.repo, - number=number, - close=args.close, - model=args.model, - grace_days=args.grace_days, - min_greptile_score=args.min_greptile_score, - ) - else: - result = triage( - repo=args.repo, - kind=kind, - number=number, - close=args.close, - model=args.model, - print_prompt=args.print_prompt, - reconsider=args.reconsider, - ) - - if result.get("action") == "print-prompt": - print(result["prompt"]) - return 0 - - summary = render_summary(result) - print(summary) - write_step_summary(summary + "\n") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/workflows/close_low_quality_prs.yml b/.github/workflows/close_low_quality_prs.yml deleted file mode 100644 index 2401be84000..00000000000 --- a/.github/workflows/close_low_quality_prs.yml +++ /dev/null @@ -1,92 +0,0 @@ -name: Close Low-Quality PRs - -# Auto-close any open PR (including drafts, regardless of age) authored by an -# external OSS contributor that Greptile reviewed with a confidence score -# below 4/5. Closures are explained in a comment that tells the contributor -# to push fixes and open a fresh PR (since OSS authors cannot reopen a PR -# closed by a bot/maintainer) or comment `@agent-shin reconsider` to have -# Agent Shin re-evaluate. -# -# Manual one-off run: -# gh workflow run "Close Low-Quality PRs" -f close=true -# -# Dry-run preview (no PRs are touched): -# gh workflow run "Close Low-Quality PRs" -f close=false - -on: - schedule: - # Daily at 09:00 UTC. Pairs well with the stale-issue workflow at midnight. - - cron: "0 9 * * *" - workflow_dispatch: - inputs: - close: - description: "Actually close matching PRs (false = dry run)." - required: false - default: "false" - type: choice - options: - - "true" - - "false" - min_age_days: - description: "Minimum PR age in days (default 0 = no age filter)." - required: false - default: "0" - min_score: - description: "Greptile score below which a PR is closed (1-5)." - required: false - default: "4" - limit: - description: "Maximum number of PRs to close in a single run." - required: false - default: "25" - -permissions: - contents: read - pull-requests: write - issues: write - -jobs: - close-low-quality-prs: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - steps: - - name: Checkout triage script - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - sparse-checkout: .github/scripts - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Run low-quality PR closer - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - # Scheduled runs are ALWAYS dry-run, even when AGENT_SHIN_ENABLED is - # "true", so the team can QA the closer's verdicts in step summaries - # before any contributor sees a PR closed. Real closures only happen - # on manual workflow_dispatch with close=true (and the variable set). - CLOSE_FLAG: ${{ github.event.inputs.close || 'false' }} - AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }} - MIN_AGE_DAYS: ${{ github.event.inputs.min_age_days || '0' }} - MIN_SCORE: ${{ github.event.inputs.min_score || '4' }} - LIMIT: ${{ github.event.inputs.limit || '25' }} - run: | - set -euo pipefail - ARGS=( - --repo "${{ github.repository }}" - --min-age-days "${MIN_AGE_DAYS}" - --min-score "${MIN_SCORE}" - --limit "${LIMIT}" - ) - if [ "${AGENT_SHIN_ENABLED:-false}" != "true" ]; then - echo "::notice::AGENT_SHIN_ENABLED is not 'true' -> forcing dry-run regardless of close input." - elif [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ] && [ "${CLOSE_FLAG}" = "true" ]; then - ARGS+=(--close) - echo "::notice::Running in close-on-fail mode." - else - echo "::notice::AGENT_SHIN_ENABLED is true but this trigger is dry-run (scheduled event or close=false)." - fi - python3 .github/scripts/close_low_quality_prs.py "${ARGS[@]}" diff --git a/.github/workflows/create_daily_oss_agent_shin_branch.yml b/.github/workflows/create_daily_oss_agent_shin_branch.yml deleted file mode 100644 index 9baf9f142f6..00000000000 --- a/.github/workflows/create_daily_oss_agent_shin_branch.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: Create Daily oss-agent-shin Branch - -on: - schedule: - - cron: "0 0 * * *" # Runs every day at midnight UTC - workflow_dispatch: # Allow manual trigger - -jobs: - create-oss-agent-shin-branch: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - permissions: - contents: write - - steps: - - name: Create daily oss-agent-shin branch - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - BRANCH_NAME="litellm_oss_agent_shin_$(date +'%m_%d_%Y')" - echo "Creating branch: $BRANCH_NAME" - if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then - echo "Branch $BRANCH_NAME already exists. Skipping creation." - exit 0 - fi - MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha') - gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent - echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA" diff --git a/.github/workflows/triage_reconsider.yml b/.github/workflows/triage_reconsider.yml deleted file mode 100644 index f35f681d09a..00000000000 --- a/.github/workflows/triage_reconsider.yml +++ /dev/null @@ -1,172 +0,0 @@ -name: Agent Shin — reconsider - -# Comment-trigger workflow: when the PR/issue author (or an internal -# collaborator) comments `@agent-shin reconsider` on a CLOSED PR/issue, -# Agent Shin re-runs LLM-judge triage on the current title+body and: -# -# - on PASS: posts a "re-evaluated and reopened" comment + reopens. -# - on FAIL: posts a "still missing X" comment and leaves it closed, -# so the contributor can iterate again. -# -# This exists because GitHub does NOT let an external (non-write-access) -# OSS contributor reopen a PR/issue closed by a bot or maintainer. Without -# this comment trigger, a contributor whose PR Agent Shin auto-closed -# would have no path back into the review queue except opening a fresh PR -# (which loses the original PR's history). The bot, on the other hand, -# has write access via GH_TOKEN and can reopen on their behalf. -# -# DRY-RUN BY DEFAULT — gated on `vars.AGENT_SHIN_ENABLED == 'true'` just -# like the other Agent Shin workflows. The workflow also gates on the -# commenter being either the PR/issue author or an internal collaborator -# (OWNER/MEMBER/COLLABORATOR) so random commenters cannot DOS the LLM -# judge or force a reopen. - -on: - issue_comment: - types: [created] - -permissions: - contents: read - issues: write - pull-requests: write - -jobs: - reconsider: - if: | - github.repository == 'BerriAI/litellm' - && contains(github.event.comment.body, '@agent-shin reconsider') - runs-on: ubuntu-latest - steps: - - name: Authorize commenter - # Only the PR/issue author OR an internal collaborator may trigger - # a reconsider. Outside random commenters could otherwise spam the - # phrase to burn LLM budget or, if a fail-open bug were ever - # introduced, force a reopen on someone else's behalf. - # - # We expose the authorization decision as a step output and gate - # every subsequent (potentially destructive) step on it. A `run:` - # step with `exit 0` would NOT stop the job — only `if:` gating - # on a known-true output is safe here. - id: auth - env: - COMMENTER: ${{ github.event.comment.user.login }} - AUTHOR: ${{ github.event.issue.user.login }} - ASSOCIATION: ${{ github.event.comment.author_association }} - run: | - set -euo pipefail - if [ "${COMMENTER}" = "${AUTHOR}" ]; then - echo "::notice::Authorized: commenter is the PR/issue author." - echo "authorized=true" >> "$GITHUB_OUTPUT" - exit 0 - fi - case "${ASSOCIATION}" in - OWNER|MEMBER|COLLABORATOR) - echo "::notice::Authorized: commenter is an internal collaborator (${ASSOCIATION})." - echo "authorized=true" >> "$GITHUB_OUTPUT" - ;; - *) - echo "::notice::Commenter '${COMMENTER}' (${ASSOCIATION}) is not authorized to trigger reconsider; skipping subsequent steps." - echo "authorized=false" >> "$GITHUB_OUTPUT" - ;; - esac - - - name: React 👀 to acknowledge the reconsider - # Add an eyes reaction to the triggering comment the moment we accept - # it, so the contributor gets instant feedback that the bot saw their - # `@agent-shin reconsider` before the slower triage steps run. Gated on - # AGENT_SHIN_ENABLED so dry-run leaves no visible trace. Best-effort: - # a reactions API hiccup must never fail the actual reconsider. - if: steps.auth.outputs.authorized == 'true' && vars.AGENT_SHIN_ENABLED == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - COMMENT_ID: ${{ github.event.comment.id }} - run: | - set -euo pipefail - gh api --method POST \ - -H "Accept: application/vnd.github+json" \ - "repos/${{ github.repository }}/issues/comments/${COMMENT_ID}/reactions" \ - -f content=eyes \ - || echo "::warning::failed to add 👀 reaction (non-fatal)" - - - name: Checkout triage script - if: steps.auth.outputs.authorized == 'true' - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - sparse-checkout: .github/scripts - persist-credentials: false - - - name: Set up Python - if: steps.auth.outputs.authorized == 'true' - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Install LLM client - if: steps.auth.outputs.authorized == 'true' - run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt - - - name: Run Agent Shin reconsider - if: steps.auth.outputs.authorized == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - # Only expose the LLM key when the bot is enabled, so a PR/issue - # author can't force paid LLM calls by spamming `@agent-shin - # reconsider` while the bot is still in dry-run. The Python script - # calls the LLM whenever this var is set (regardless of `--close`); - # stripping `--close` doesn't suppress the API call, only the - # destructive side effects. Mirror the gating used by every other - # Agent Shin workflow (triage_pr_with_llm.yml, review_gate.yml, ...). - OPENAI_API_KEY: ${{ vars.AGENT_SHIN_ENABLED == 'true' && secrets.OPENAI_API_KEY || '' }} - OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }} - TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }} - AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }} - # `issue_comment` events fire for both issues and PR comments. - # `issue.pull_request` is set iff this is a PR comment, so we use - # its presence to decide whether to invoke `--pr N` or `--issue N`. - IS_PR: ${{ github.event.issue.pull_request != null }} - NUMBER: ${{ github.event.issue.number }} - run: | - set -euo pipefail - if [ "${IS_PR}" = "true" ]; then - ARGS=(--repo "${{ github.repository }}" --pr "${NUMBER}" --reconsider) - else - ARGS=(--repo "${{ github.repository }}" --issue "${NUMBER}" --reconsider) - fi - # Reconsider's destructive actions (post comment + reopen) are - # gated on `--close`, mirroring the regular triage workflows. - # When AGENT_SHIN_ENABLED is not the EXACT string "true", we - # still run the script so its verdict + would-X action lands in - # the step summary for QA — but without `--close`, the script - # returns `would-reopen` / `would-reconsider-still-failing` - # instead of touching GitHub state. - # - # Use the positive `= "true"` gate (not `!= "true" -> exit`) so - # the workflow guardrails in - # tests/test_litellm/test_github_triage_workflows.py see the - # canonical fail-safe enable pattern. Unknown values like - # "True", "yes", "1", or typos fall through to the dry-run - # branch, which is the safe default. - if [ "${AGENT_SHIN_ENABLED:-false}" = "true" ]; then - ARGS+=(--close) - echo "::notice::Agent Shin reconsider ENABLED — running real triage (close=true)." - else - echo "::notice::AGENT_SHIN_ENABLED is not 'true' -> reconsider stays in dry-run (no comment, no reopen)." - fi - python3 .github/scripts/triage_with_llm.py "${ARGS[@]}" - - - name: React 👍 when the reconsider finishes - # Once the reconsider run has completed successfully, add a thumbs-up so - # the contributor sees the bot is done (the 👀 stays, signalling - # seen -> handled). `success()` keeps this from firing if the run - # errored, and the AGENT_SHIN_ENABLED gate keeps dry-run inert. - if: success() && steps.auth.outputs.authorized == 'true' && vars.AGENT_SHIN_ENABLED == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - COMMENT_ID: ${{ github.event.comment.id }} - run: | - set -euo pipefail - gh api --method POST \ - -H "Accept: application/vnd.github+json" \ - "repos/${{ github.repository }}/issues/comments/${COMMENT_ID}/reactions" \ - -f content=+1 \ - || echo "::warning::failed to add 👍 reaction (non-fatal)" diff --git a/tests/test_litellm/test_github_close_low_quality_prs.py b/tests/test_litellm/test_github_close_low_quality_prs.py deleted file mode 100644 index 2a891ca72f5..00000000000 --- a/tests/test_litellm/test_github_close_low_quality_prs.py +++ /dev/null @@ -1,856 +0,0 @@ -"""Unit tests for `.github/scripts/close_low_quality_prs.py`. - -These exercise the pure logic (score extraction and per-PR evaluation) without -hitting GitHub. Network/CLI calls are stubbed via monkeypatch. -""" - -from __future__ import annotations - -import datetime as dt -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] - / ".github" - / "scripts" - / "close_low_quality_prs.py" -) - - -@pytest.fixture(scope="module") -def closer_module(): - """Load the script as a module via its file path (it lives outside the package).""" - spec = importlib.util.spec_from_file_location("close_low_quality_prs", SCRIPT_PATH) - assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" - module = importlib.util.module_from_spec(spec) - sys.modules["close_low_quality_prs"] = module - spec.loader.exec_module(module) - return module - - -def _greptile_comment( - body: str, - updated_at: str = "2026-05-10T00:00:00Z", - login: str = "greptile-apps[bot]", -) -> dict: - return { - "user": {"login": login}, - "body": body, - "created_at": updated_at, - "updated_at": updated_at, - } - - -class TestExtractGreptileScore: - def test_should_extract_score_from_html_header(self, closer_module): - comments = [ - _greptile_comment("

Confidence Score: 3/5

\nSome body text.") - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 3 - - def test_should_accept_both_greptile_login_variants(self, closer_module): - # REST API form ("greptile-apps[bot]") and GraphQL form ("greptile-apps") - for login in ("greptile-apps", "greptile-apps[bot]"): - comments = [ - _greptile_comment("

Confidence Score: 2/5

", login=login) - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None, f"failed to detect score for login={login}" - score, _ = result - assert score == 2 - - def test_should_extract_score_from_plain_text(self, closer_module): - comments = [_greptile_comment("Confidence Score: 5/5 — looks good!")] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 5 - - def test_should_tolerate_whitespace_and_case(self, closer_module): - comments = [_greptile_comment("**confidence score : 2 / 5**")] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 2 - - def test_should_pick_most_recent_comment_when_rereview_happens(self, closer_module): - comments = [ - _greptile_comment( - "Confidence Score: 2/5", updated_at="2026-05-01T00:00:00Z" - ), - _greptile_comment( - "Confidence Score: 5/5", updated_at="2026-05-12T00:00:00Z" - ), - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 5 - - def test_should_ignore_non_greptile_authors(self, closer_module): - comments = [ - { - "user": {"login": "some-human"}, - "body": "Confidence Score: 1/5", - "created_at": "2026-05-12T00:00:00Z", - "updated_at": "2026-05-12T00:00:00Z", - } - ] - assert closer_module.extract_greptile_score(comments) is None - - def test_should_return_none_when_no_score_present(self, closer_module): - comments = [_greptile_comment("Greptile summary without a score.")] - assert closer_module.extract_greptile_score(comments) is None - - def test_should_return_none_for_empty_comments(self, closer_module): - assert closer_module.extract_greptile_score([]) is None - - -class TestEvaluatePr: - @pytest.fixture(autouse=True) - def _now(self): - return dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - - def _make_pr( - self, - *, - number: int = 1, - created_days_ago: int = 10, - is_draft: bool = False, - labels: list[str] | None = None, - author_login: str = "mateo-berri", - ) -> dict: - created = dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - dt.timedelta( - days=created_days_ago - ) - return { - "number": number, - "title": f"PR #{number}", - "createdAt": created.isoformat().replace("+00:00", "Z"), - "isDraft": is_draft, - "labels": [{"name": lbl} for lbl in (labels or [])], - "author": {"login": author_login}, - "url": f"https://example.com/pr/{number}", - } - - @pytest.fixture(autouse=True) - def _external_author(self, closer_module, monkeypatch): - """Treat every test PR as external unless overridden.""" - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: True - ) - - def test_should_warn_drafts_when_score_low_first_time( - self, closer_module, _now, monkeypatch - ): - # Drafts are NOT a free pass — the open-PR queue should reflect any - # PR that needs human attention regardless of draft status. Authors - # who need a long-lived draft can use the `wip` opt-out label. - # First run: warn the contributor (1-day grace), don't close yet. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 2/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(is_draft=True, created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 2 and age == 0 - - def test_should_warn_brand_new_pr_when_min_age_zero( - self, closer_module, _now, monkeypatch - ): - # `min_age_days=0` means no age filter — a freshly-opened PR is - # eligible the moment Greptile scores it below threshold. The - # first detection still goes through the warn-grace step rather - # than closing immediately, giving the contributor 2 hours to - # respond before the next run actually closes the PR. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 1/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 1 and age == 0 - - def test_should_skip_optout_label_case_insensitive( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for opt-outs"), - ) - action, _, _ = closer_module.evaluate_pr( - self._make_pr(labels=["WIP"]), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels={"wip"}, - ) - assert action == "skip-optout-label" - - def test_should_skip_too_young_when_min_age_set( - self, closer_module, _now, monkeypatch - ): - # The min-age-days flag is now opt-in (default 0). When a maintainer - # explicitly passes a positive value (e.g. for a backfill run that - # wants to spare brand-new PRs), the skip-too-young path still works. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for young PRs"), - ) - action, _, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=2), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-too-young" - assert age == 2 - - def test_should_not_skip_when_min_age_is_zero( - self, closer_module, _now, monkeypatch - ): - # With the new default min_age_days=0, even a 0-day-old PR is - # evaluated. This test pins that behavior so future refactors don't - # silently restore an age filter. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 5/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-score-ok" - assert score == 5 and age == 0 - - def test_should_skip_when_greptile_has_not_reviewed( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr(closer_module, "fetch_pr_comments", lambda *a, **kw: []) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-no-greptile-score" - assert score is None and age == 10 - - def test_should_skip_when_score_meets_threshold( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 4/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-score-ok" - assert score == 4 and age == 10 - - def test_should_warn_when_old_and_low_score_no_prior_warning( - self, closer_module, _now, monkeypatch - ): - # Even an old PR that still has no grace warning gets one on the - # first eligible run — the daily cron is the natural cadence, so - # an existing-but-never-warned PR enters the grace flow normally. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 3/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 3 and age == 10 - - def test_should_close_when_grace_warning_aged_out_and_score_still_low( - self, closer_module, _now, monkeypatch - ): - # Day-1 the closer posted a warning. Day-2 the PR still scores <4 - # AND the warning is older than `GRACE_PERIOD_SECONDS`, so the - # action flips to `close`. This is the "grace expired" path. - old_warning = { - "user": {"login": "github-actions[bot]"}, - "body": ( - "you have 2 hours to fix this\n\n" + closer_module.GRACE_COMMENT_MARKER - ), - "created_at": ( - _now - dt.timedelta(seconds=closer_module.GRACE_PERIOD_SECONDS + 60) - ) - .isoformat() - .replace("+00:00", "Z"), - "updated_at": "2026-05-15T00:00:00Z", - } - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment( - "

Confidence Score: 1/5

", - updated_at="2026-05-15T00:00:00Z", - ), - old_warning, - ], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=14), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "close" - assert score == 1 - - def test_should_skip_when_grace_warning_within_window( - self, closer_module, _now, monkeypatch - ): - # Within the 2-hour grace window the closer must NOT close the - # PR even if the score is still low. The warning is only an hour - # old; give the contributor time to push fixes before destruction. - recent_warning = { - "user": {"login": "github-actions[bot]"}, - "body": "warning text\n\n" + closer_module.GRACE_COMMENT_MARKER, - "created_at": (_now - dt.timedelta(hours=1)) - .isoformat() - .replace("+00:00", "Z"), - } - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment("Confidence Score: 2/5"), - recent_warning, - ], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-in-grace-period" - assert score == 2 - - def test_should_warn_grace_for_swiftwinds_not_close_immediately( - self, closer_module, _now, monkeypatch - ): - # Regression: SwiftWinds (the dogfood account) used to be in a - # now-removed `IMMEDIATE_CLOSE_LOGINS` bypass that closed on first - # detection. It must now follow the SAME grace path as every other - # external author: warn first, close only after the window elapses. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 1/5")], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0, author_login="SwiftWinds"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 1 - - def test_should_skip_internal_authors(self, closer_module, _now, monkeypatch): - # Override the fixture for this one test. - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for internal"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=14, author_login="krrishdholakia"), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - allowlist=frozenset(), - ) - assert action == "skip-internal" - assert score is None - - -class TestMainOptoutLabelDefault: - """`--optout-label` must REPLACE the canonical defaults, not append.""" - - def _patch_no_op(self, closer_module, monkeypatch): - monkeypatch.setattr(closer_module, "fetch_open_prs", lambda repo: []) - # `optout_labels` is captured indirectly via evaluate_pr; sniff the - # set passed in by stubbing evaluate_pr. - captured: dict = {} - - def fake_evaluate(pr, now, min_age_days, min_score, repo, optout_labels): - captured["optout_labels"] = set(optout_labels) - return ("skip-internal", None, None) - - monkeypatch.setattr(closer_module, "evaluate_pr", fake_evaluate) - return captured - - def test_should_use_canonical_defaults_when_flag_omitted( - self, closer_module, monkeypatch - ): - captured = self._patch_no_op(closer_module, monkeypatch) - # No PRs -> capture won't fire; instead inject one synthetic PR via - # fetch_open_prs so evaluate_pr is invoked at least once. - monkeypatch.setattr( - closer_module, - "fetch_open_prs", - lambda repo: [ - { - "number": 1, - "title": "p", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": True, - "labels": [], - "author": {"login": "x"}, - } - ], - ) - monkeypatch.setattr(sys, "argv", ["close_low_quality_prs.py"]) - rc = closer_module.main() - assert rc == 0 - assert captured["optout_labels"] == set(closer_module.DEFAULT_OPTOUT_LABELS) - - def test_should_replace_defaults_when_flag_provided( - self, closer_module, monkeypatch - ): - captured = self._patch_no_op(closer_module, monkeypatch) - monkeypatch.setattr( - closer_module, - "fetch_open_prs", - lambda repo: [ - { - "number": 1, - "title": "p", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": True, - "labels": [], - "author": {"login": "x"}, - } - ], - ) - monkeypatch.setattr( - sys, - "argv", - [ - "close_low_quality_prs.py", - "--optout-label", - "hold", - "--optout-label", - "needs-discussion", - ], - ) - rc = closer_module.main() - assert rc == 0 - # Crucially, none of the canonical defaults leak in. - assert captured["optout_labels"] == {"hold", "needs-discussion"} - for default in closer_module.DEFAULT_OPTOUT_LABELS: - assert default not in captured["optout_labels"], default - - -class TestSecondsSinceLastGraceWarning: - """Grace-period detection: only counts comments by the bot identity - that contain the shared `GRACE_COMMENT_MARKER`.""" - - def _make_marker_comment( - self, - closer_module, - *, - login: str = "github-actions[bot]", - created_at: str = "2026-05-16T00:00:00Z", - include_marker: bool = True, - ) -> dict: - body = "warning text" - if include_marker: - body += "\n\n" + closer_module.GRACE_COMMENT_MARKER - return { - "user": {"login": login}, - "body": body, - "created_at": created_at, - } - - def test_should_return_none_when_no_marker_comment(self, closer_module): - comments = [ - { - "user": {"login": "github-actions[bot]"}, - "body": "Some other bot comment", - "created_at": "2026-05-16T00:00:00Z", - } - ] - assert closer_module.seconds_since_last_grace_warning(comments) is None - - def test_should_return_none_for_empty(self, closer_module): - assert closer_module.seconds_since_last_grace_warning([]) is None - - def test_should_ignore_non_bot_comments_with_marker(self, closer_module): - # If a curious user quotes the marker in a comment, we must NOT - # treat it as a bot warning. The grace timer would then never fire. - comments = [ - self._make_marker_comment(closer_module, login="random-user"), - ] - assert closer_module.seconds_since_last_grace_warning(comments) is None - - def test_should_pick_latest_marker_comment(self, closer_module): - # When multiple grace warnings exist (e.g. a re-open cycle), use - # the most recent one to compute the age. - comments = [ - self._make_marker_comment(closer_module, created_at="2026-05-15T00:00:00Z"), - self._make_marker_comment(closer_module, created_at="2026-05-16T23:00:00Z"), - ] - now = dt.datetime(2026, 5, 17, 0, 0, 0, tzinfo=dt.timezone.utc) - age = closer_module.seconds_since_last_grace_warning(comments, now=now) - # 1h = 3600s - assert age == 3600.0 - - -class TestGraceWarningCommentText: - """Pin the user-facing language in the grace warning comment so the - grace-window and `@greptileai still works after close` promises - don't get accidentally dropped in a future refactor. - """ - - def test_should_state_grace_window(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - # The user's PR explicitly said "specify in the comment" — pin - # that the grace window appears in the comment. - assert "2 hours" in body - - def test_should_mention_agent_shin_reconsider(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert "@agent-shin reconsider" in body - - def test_should_promise_greptileai_works_after_close(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_should_carry_grace_marker(self, closer_module): - # The marker is what `seconds_since_last_grace_warning` greps for - # to detect a prior warning — dropping it would silently break - # the cooldown. - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert closer_module.GRACE_COMMENT_MARKER in body - - def test_close_comment_should_mention_greptileai_post_close(self, closer_module): - # The close comment should ALSO point at the @greptileai post-close - # re-review path so contributors see the same options whether they - # read the warning or only catch the close comment. - body = closer_module.format_close_comment(score=2, threshold=4) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_close_comment_should_advertise_reconsider(self, closer_module): - body = closer_module.format_close_comment(score=2, threshold=4) - assert "@agent-shin reconsider" in body - - def test_close_comment_should_carry_agent_shin_close_marker(self, closer_module): - # The close comment advertises `@agent-shin reconsider`, and the - # reconsider reopen guard (`was_closed_by_agent_shin`) only treats a - # PR as Agent-Shin-closed when the close comment carries this marker. - # Dropping it silently breaks the advertised recovery path for every - # PR closed by this daily sweep. - body = closer_module.format_close_comment(score=2, threshold=4) - assert closer_module.AGENT_SHIN_CLOSE_MARKER in body - - def test_close_comment_should_state_score_and_threshold(self, closer_module): - body = closer_module.format_close_comment(score=1, threshold=4) - assert "1/5" in body - assert "4/5" in body - - -class TestHasOptoutLabel: - def test_should_match_label_case_insensitively(self, closer_module): - pr = {"labels": [{"name": "Do Not Close"}, {"name": "bug"}]} - assert closer_module.has_optout_label(pr, {"do not close"}) is True - - def test_should_return_false_when_no_match(self, closer_module): - pr = {"labels": [{"name": "bug"}, {"name": "enhancement"}]} - assert closer_module.has_optout_label(pr, {"wip", "keep open"}) is False - - def test_should_handle_missing_labels(self, closer_module): - assert closer_module.has_optout_label({}, {"wip"}) is False - - -class TestListOpenItemsNoCap: - """The bulk sweeps must fetch the ENTIRE open backlog. - - Regression guard for the old hard-coded ``--limit 1000``: gh lists - newest-first, so a low cap silently dropped the *oldest* PRs/issues — - exactly the stale ones a low-quality sweep exists to catch. - """ - - @staticmethod - def _shared(closer_module): - # `closer_module` loading puts `.github/scripts` on sys.path and - # imports agent_shin_shared, so it's already in sys.modules. - import agent_shin_shared - - return agent_shin_shared - - def _capture_gh_args(self, closer_module, monkeypatch, *, returns="[]"): - shared = self._shared(closer_module) - captured: dict = {} - - def fake_gh(*args): - captured["args"] = args - return returns - - # `list_open_items` looks up `gh` in agent_shin_shared's namespace. - monkeypatch.setattr(shared, "gh", fake_gh) - return shared, captured - - def test_list_open_items_passes_no_cap_limit_not_1000( - self, closer_module, monkeypatch - ): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("pr", repo="o/r", fields="number,title") - args = captured["args"] - assert "--limit" in args - limit_value = args[args.index("--limit") + 1] - assert limit_value == str(shared.GH_LIST_ALL_LIMIT) - assert limit_value != "1000" - # A meaningful ceiling: comfortably above any realistic open backlog. - assert shared.GH_LIST_ALL_LIMIT >= 100_000 - - def test_list_open_items_uses_dedicated_command_state_and_fields( - self, closer_module, monkeypatch - ): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("issue", repo="o/r", fields="number") - args = captured["args"] - assert args[0] == "issue" and args[1] == "list" - assert args[args.index("--state") + 1] == "open" - assert args[args.index("--json") + 1] == "number" - assert tuple(args[-2:]) == ("--repo", "o/r") - - def test_list_open_items_omits_repo_when_none(self, closer_module, monkeypatch): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("pr", repo=None, fields="number") - assert "--repo" not in captured["args"] - - def test_list_open_items_parses_json_array(self, closer_module, monkeypatch): - shared, _ = self._capture_gh_args( - closer_module, monkeypatch, returns='[{"number": 1}, {"number": 2}]' - ) - items = shared.list_open_items("pr", repo=None, fields="number") - assert [i["number"] for i in items] == [1, 2] - - def test_list_open_items_rejects_unknown_kind(self, closer_module): - shared = self._shared(closer_module) - with pytest.raises(ValueError, match="kind must be 'pr' or 'issue', got 'both"): - shared.list_open_items("both", repo="o/r", fields="number") - - def test_fetch_open_prs_delegates_with_no_cap(self, closer_module, monkeypatch): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - closer_module.fetch_open_prs("o/r") - args = captured["args"] - assert args[0] == "pr" - assert args[args.index("--limit") + 1] == str(shared.GH_LIST_ALL_LIMIT) - # Still requests every field downstream evaluate_pr / labels logic needs. - assert "createdAt" in args[args.index("--json") + 1] - - -class TestEvaluatePrAllowlist: - """While the dogfood allowlist is active `evaluate_pr` only acts on the - named accounts and bypasses the external-only restriction for them. - Emptying it restores the internal-author skip.""" - - @pytest.fixture(autouse=True) - def _now(self): - return dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - - def _make_pr(self, *, author_login: str, created_days_ago: int = 10) -> dict: - created = dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - dt.timedelta( - days=created_days_ago - ) - return { - "number": 1, - "title": "PR #1", - "createdAt": created.isoformat().replace("+00:00", "Z"), - "isDraft": False, - "labels": [], - "author": {"login": author_login}, - "url": "https://example.com/pr/1", - } - - def test_should_skip_author_not_on_allowlist( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("must not fetch comments for non-allowlisted"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="random-oss-dev"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-not-allowlisted" - assert score is None - - def test_should_act_on_allowlisted_internal_author( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 2/5")], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="mateo-berri", created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 2 - - def test_empty_allowlist_restores_internal_skip( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("must not fetch comments for internal"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="krrishdholakia"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - allowlist=frozenset(), - ) - assert action == "skip-internal" - - def test_allowlist_constant_is_the_two_dogfood_accounts(self, closer_module): - assert closer_module.ALLOWLIST_LOGINS == frozenset( - {"mateo-berri", "swiftwinds"} - ) - - -class TestDryRunGateOnClose: - """Regression: the daily sweep is dry-run unless `--close` is passed - (the workflow only adds it when `AGENT_SHIN_ENABLED=true`). A closeable - PR (low score, grace window elapsed) must be DETECTED and reported as - "would close", but the dry run must never make a real GitHub mutation, - so merging Agent Shin stays inert by default.""" - - def _closeable_pr(self) -> dict: - return { - "number": 7, - "title": "thin PR", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": False, - "labels": [], - "author": {"login": "SwiftWinds"}, - "url": "https://example.com/pr/7", - } - - def test_dry_run_sweep_detects_but_does_not_close( - self, closer_module, monkeypatch, capsys - ): - aged_out_warning = { - "user": {"login": "github-actions[bot]"}, - "body": "warned\n\n" + closer_module.GRACE_COMMENT_MARKER, - # Far enough in the past that it's aged out regardless of - # GRACE_PERIOD_SECONDS, since main() pins `now` to real time. - "created_at": "2020-01-01T00:00:00Z", - } - monkeypatch.setattr( - closer_module, "fetch_open_prs", lambda repo: [self._closeable_pr()] - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment("Confidence Score: 1/5"), - aged_out_warning, - ], - ) - # Any real GitHub mutation during a dry run is the bug under test. - monkeypatch.setattr( - closer_module, - "gh", - lambda *a, **kw: pytest.fail(f"dry run must not call gh: {a}"), - ) - monkeypatch.setattr(sys, "argv", ["close_low_quality_prs.py"]) - - rc = closer_module.main() - - assert rc == 0 - # The PR is detected as closeable, just not acted on. - assert "Total would close: 1" in capsys.readouterr().out diff --git a/tests/test_litellm/test_github_review_gate.py b/tests/test_litellm/test_github_review_gate.py deleted file mode 100644 index 001fa8f43f5..00000000000 --- a/tests/test_litellm/test_github_review_gate.py +++ /dev/null @@ -1,524 +0,0 @@ -"""Unit tests for the `ready for review` label lifecycle (Agent Shin review gate). - -Exercises `triage_with_llm.review_gate`, the state machine that keeps the -`ready for review` label in sync with whether a PR clears both the LLM rubric -and Greptile's confidence score: - - * pass (untagged) -> add label + "ready for review" comment - * pass (untagged, recovered) -> add label + "all clear again" comment - * pass (already tagged) -> noop - * regress (tagged) -> remove label + "what's missing" comment, stays open - * fail (untagged, within 24h)-> one-time "what's missing" notice - * fail (untagged, >24h) -> close + comment - * dry run (close=False) -> would-* previews, no side effects -""" - -from __future__ import annotations - -import datetime as dt -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] / ".github" / "scripts" / "triage_with_llm.py" -) - -NOW = dt.datetime(2026, 5, 24, 12, 0, 0, tzinfo=dt.timezone.utc) -JUST_NOW = "2026-05-24T11:00:00Z" # 1h old -> within 24h grace -TWO_DAYS_AGO = "2026-05-22T11:00:00Z" # >24h old -> past grace - - -@pytest.fixture(scope="module") -def triage_module(): - spec = importlib.util.spec_from_file_location("triage_with_llm", SCRIPT_PATH) - assert spec and spec.loader - module = importlib.util.module_from_spec(spec) - sys.modules["triage_with_llm"] = module - spec.loader.exec_module(module) - return module - - -class _Recorder: - """Captures every gh mutation review_gate could fire, and fails loudly - on the ones a given scenario forbids.""" - - def __init__(self, triage_module, monkeypatch): - self.comments: list[str] = [] - self.added: list[str] = [] - self.removed: list[str] = [] - self.closed: list[int] = [] - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: self.comments.append(body), - ) - monkeypatch.setattr( - triage_module, - "add_label", - lambda repo, n, label: self.added.append(label), - ) - monkeypatch.setattr( - triage_module, - "remove_label", - lambda repo, n, label: self.removed.append(label), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda repo, n: self.closed.append(n), - ) - - -def _make_pr(**overrides): - base = { - "number": 7, - "title": "feat: do a thing", - "body": "some body without a linked issue or QA proof", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - "labels": [], - "created_at": JUST_NOW, - } - base.update(overrides) - return base - - -def _pass(prompt): - return '{"verdict": "pass", "missing": [], "explanation": "looks good"}' - - -def _fail(prompt): - return ( - '{"verdict": "fail", "missing": ["QA proof", "expected vs. actual"],' - ' "explanation": "thin description"}' - ) - - -def _gate(triage_module, **kwargs): - """Call review_gate with safe defaults for the injectable hooks.""" - params = dict( - repo="o/r", - number=7, - close=True, - model="m", - judge=_pass, - greptile_score=None, - comments=[], - now=NOW, - ) - params.update(kwargs) - return triage_module.review_gate(**params) - - -class TestReviewGatePass: - def test_pass_untagged_adds_label_and_ready_comment( - self, triage_module, monkeypatch - ): - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=5) - - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - assert rec.removed == [] and rec.closed == [] - assert len(rec.comments) == 1 - assert "ready for review" in rec.comments[0].lower() - assert triage_module.READY_MARKER in rec.comments[0] - assert "5/5" in rec.comments[0] - - def test_pass_already_tagged_is_noop(self, triage_module, monkeypatch): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=5) - - assert result["action"] == "noop-passing" - assert rec.added == [] and rec.removed == [] and rec.comments == [] - - def test_pass_after_prior_regression_uses_all_clear_wording( - self, triage_module, monkeypatch - ): - # A regression marker in history -> this is a recovery, not a first pass. - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - } - ] - - result = _gate(triage_module, judge=_pass, greptile_score=5, comments=prior) - - assert result["action"] == "labeled-ready" - assert "all clear" in rec.comments[0].lower() - - def test_linked_issue_passes_without_calling_judge( - self, triage_module, monkeypatch - ): - pr = _make_pr(body="Fixes #4321\n\nbody") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate( - triage_module, - judge=lambda p: pytest.fail("LLM must not be called for linked issue"), - greptile_score=5, - ) - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - - -class TestReviewGateRegression: - def test_regression_removes_label_and_keeps_pr_open( - self, triage_module, monkeypatch - ): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=5) - - assert result["action"] == "label-removed-regressed" - assert rec.removed == [triage_module.READY_FOR_REVIEW_LABEL] - assert rec.closed == [] # regression NEVER closes the PR - assert triage_module.REGRESSED_MARKER in rec.comments[0] - assert "QA proof" in rec.comments[0] - # The state machine closes a still-failing PR `grace_days` after this - # notice (default 24h); the comment must disclose that deadline rather - # than implying the PR stays open indefinitely. - assert "24 hours" in rec.comments[0] - assert "auto-closed" in rec.comments[0] - - def test_regression_comment_discloses_grace_deadline(self, triage_module): - one_day = triage_module.format_regression_comment( - ["QA proof"], "needs work", grace_days=1 - ) - assert "24 hours" in one_day - assert "auto-closed" in one_day - - three_days = triage_module.format_regression_comment( - ["QA proof"], "needs work", grace_days=3 - ) - assert "3 days" in three_days - assert "auto-closed" in three_days - - def test_greptile_drop_alone_triggers_regression(self, triage_module, monkeypatch): - # Rubric still passes, but Greptile fell to 2/5 -> not passing. - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=2) - - assert result["action"] == "label-removed-regressed" - assert rec.removed == [triage_module.READY_FOR_REVIEW_LABEL] - assert "2/5" in rec.comments[0] - - def test_greptile_score_read_from_comments_when_not_injected( - self, triage_module, monkeypatch - ): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - greptile = [ - { - "user": {"login": "greptile-apps[bot]"}, - "body": "Confidence Score: 2/5", - "created_at": "2026-05-24T10:00:00Z", - } - ] - - result = _gate( - triage_module, - judge=_pass, - greptile_score=triage_module._UNSET, - comments=greptile, - ) - assert result["action"] == "label-removed-regressed" - assert "2/5" in rec.comments[0] - - -class TestReviewGateGraceAndClose: - def test_within_grace_posts_one_time_notice(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, "fetch_pr", lambda repo, n: _make_pr(created_at=JUST_NOW) - ) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=None) - - assert result["action"] == "within-grace-notified" - assert rec.closed == [] and rec.added == [] and rec.removed == [] - assert triage_module.WITHIN_GRACE_MARKER in rec.comments[0] - assert "QA proof" in rec.comments[0] - - def test_within_grace_does_not_double_notify(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, "fetch_pr", lambda repo, n: _make_pr(created_at=JUST_NOW) - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.WITHIN_GRACE_MARKER, - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "within-grace-already-notified" - assert rec.comments == [] - - def test_past_grace_closes_with_comment(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=None) - - assert result["action"] == "closed" - assert rec.closed == [7] - assert len(rec.comments) == 1 - # The close comment must carry the reconsider provenance marker so - # `was_closed_by_agent_shin` can later recognize this as an Agent Shin - # close (and not some other workflow's `github-actions[bot]` close). - assert triage_module.AGENT_SHIN_CLOSE_MARKER in rec.comments[0] - - def test_recent_regression_marker_blocks_close(self, triage_module, monkeypatch): - """A failing PR with a fresh regression notice must NOT be closed — - the contributor needs a window to address the regression.""" - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - # Posted just an hour before NOW -> well inside grace_days. - "created_at": "2026-05-24T11:00:00Z", - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "regressed-already-notified" - assert rec.closed == [] and rec.comments == [] - - def test_stale_regression_marker_allows_close(self, triage_module, monkeypatch): - """Once grace_days have elapsed since the regression notice, the - review gate must let the close path fire — otherwise PRs that were - regressed and then abandoned stay open forever.""" - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - # Posted 30 days before NOW -> well past the default 1-day grace. - "created_at": "2026-04-24T11:00:00Z", - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "closed" - assert rec.closed == [7] - assert len(rec.comments) == 1 - - def test_linked_issue_with_greptile_fail_uses_greptile_explanation( - self, triage_module, monkeypatch - ): - """When the rubric short-circuits to pass (linked-issue regex) but - Greptile dragged the PR under the bar, the close comment's - explanation must describe the Greptile shortfall, not the - misleading "LLM was not called" rubric placeholder.""" - pr = _make_pr(body="Fixes #4321\n\nbody", created_at=TWO_DAYS_AGO) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate( - triage_module, - judge=lambda p: pytest.fail("LLM must not be called for linked issue"), - greptile_score=2, - ) - - assert result["action"] == "closed" - assert len(rec.comments) == 1 - body = rec.comments[0] - assert "LLM was not called" not in body - assert "Greptile" in body and "2/5" in body - - -class TestReviewGateDryRun: - @pytest.mark.parametrize( - "scenario,labels,judge,score,created,expected", - [ - ("pass", [], _pass, 5, JUST_NOW, "would-label-ready"), - ( - "regress", - [{"name": "ready for review"}], - _fail, - 5, - JUST_NOW, - "would-remove-label", - ), - ("within-grace", [], _fail, None, JUST_NOW, "would-notify-within-grace"), - ("past-grace", [], _fail, None, TWO_DAYS_AGO, "would-close"), - ], - ) - def test_dry_run_previews_without_side_effects( - self, - triage_module, - monkeypatch, - scenario, - labels, - judge, - score, - created, - expected, - ): - pr = _make_pr(labels=labels, created_at=created) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, close=False, judge=judge, greptile_score=score) - - assert result["action"] == expected - # Dry run touches nothing. - assert rec.added == [] and rec.removed == [] and rec.closed == [] - assert rec.comments == [] - assert "comment" in result # preview body still surfaced - - -class TestReviewGateGuards: - def test_skips_internal_author(self, triage_module, monkeypatch): - pr = _make_pr(author_association="MEMBER", user={"login": "krrish"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate( - triage_module, - judge=lambda p: pytest.fail("no LLM for internal"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_skips_closed_pr(self, triage_module, monkeypatch): - pr = _make_pr(state="closed") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate(triage_module, judge=lambda p: pytest.fail("no LLM for closed")) - assert result["action"] == "skip-not-open" - - def test_llm_error_is_non_destructive(self, triage_module, monkeypatch): - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - - def boom(prompt): - raise RuntimeError("api down") - - result = _gate(triage_module, judge=boom, greptile_score=None) - - assert result["action"] == "skip-llm-error" - assert rec.closed == [] and rec.added == [] and rec.removed == [] - - def test_full_recovery_cycle(self, triage_module, monkeypatch): - """pass -> regress -> recover, threading labels/comments like GitHub would.""" - state = {"labels": [], "comments": []} - - def fake_fetch(repo, n): - return _make_pr(labels=list(state["labels"]), created_at=JUST_NOW) - - monkeypatch.setattr(triage_module, "fetch_pr", fake_fetch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: state["comments"].append( - {"user": {"login": "github-actions[bot]"}, "body": body} - ), - ) - monkeypatch.setattr( - triage_module, - "add_label", - lambda repo, n, label: state["labels"].append({"name": label}), - ) - monkeypatch.setattr( - triage_module, - "remove_label", - lambda repo, n, label: state["labels"].clear(), - ) - monkeypatch.setattr( - triage_module, "close_pr", lambda repo, n: pytest.fail("must not close") - ) - - # 1) passes -> tagged - r1 = _gate( - triage_module, judge=_pass, greptile_score=5, comments=state["comments"] - ) - assert r1["action"] == "labeled-ready" - assert any(lbl["name"] == "ready for review" for lbl in state["labels"]) - - # 2) regresses -> tag removed, comment posted, PR still open - r2 = _gate( - triage_module, judge=_fail, greptile_score=2, comments=state["comments"] - ) - assert r2["action"] == "label-removed-regressed" - assert state["labels"] == [] - - # 3) fixed again -> "all clear" + tag back - r3 = _gate( - triage_module, judge=_pass, greptile_score=5, comments=state["comments"] - ) - assert r3["action"] == "labeled-ready" - assert any(lbl["name"] == "ready for review" for lbl in state["labels"]) - assert "all clear" in state["comments"][-1]["body"].lower() - - -class TestReviewGateAllowlist: - """While the dogfood allowlist is active it is the sole author gate: - only the named accounts pass, and for them the internal-author exemption - is bypassed. Emptying it restores the normal internal-author skip.""" - - def test_should_skip_author_not_on_allowlist(self, triage_module, monkeypatch): - pr = _make_pr(user={"login": "random-oss-dev"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - result = _gate( - triage_module, judge=lambda p: pytest.fail("no LLM for non-allowlisted") - ) - assert result["action"] == "skip-not-allowlisted" - assert rec.added == [] and rec.comments == [] and rec.closed == [] - - def test_should_act_on_allowlisted_internal_author( - self, triage_module, monkeypatch - ): - pr = _make_pr(author_association="MEMBER", user={"login": "mateo-berri"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - result = _gate(triage_module, judge=_pass, greptile_score=5) - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - - def test_empty_allowlist_restores_internal_skip(self, triage_module, monkeypatch): - pr = _make_pr(author_association="MEMBER", user={"login": "krrish"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate( - triage_module, - judge=lambda p: pytest.fail("no LLM for internal"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" diff --git a/tests/test_litellm/test_github_triage_with_llm.py b/tests/test_litellm/test_github_triage_with_llm.py deleted file mode 100644 index ddffb978b48..00000000000 --- a/tests/test_litellm/test_github_triage_with_llm.py +++ /dev/null @@ -1,2134 +0,0 @@ -"""Unit tests for `.github/scripts/triage_with_llm.py` (Agent Shin).""" - -from __future__ import annotations - -import importlib.util -import json -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] / ".github" / "scripts" / "triage_with_llm.py" -) - - -@pytest.fixture(scope="module") -def triage_module(): - spec = importlib.util.spec_from_file_location("triage_with_llm", SCRIPT_PATH) - assert spec and spec.loader - module = importlib.util.module_from_spec(spec) - sys.modules["triage_with_llm"] = module - spec.loader.exec_module(module) - return module - - -class TestIsInternalContributor: - @pytest.mark.parametrize("association", ["OWNER", "MEMBER", "COLLABORATOR"]) - def test_should_mark_org_associations_as_internal(self, triage_module, association): - item = { - "author_association": association, - "user": {"login": "krrishdholakia"}, - } - assert triage_module.is_internal_contributor(item) is True - - @pytest.mark.parametrize( - "association", - ["CONTRIBUTOR", "FIRST_TIME_CONTRIBUTOR", "FIRST_TIMER", "NONE"], - ) - def test_should_mark_outside_associations_as_external( - self, triage_module, association - ): - item = { - "author_association": association, - "user": {"login": "random-oss-dev"}, - } - assert triage_module.is_internal_contributor(item) is False - - @pytest.mark.parametrize( - "item", - [ - {"author_association": "", "user": {"login": "random-oss-dev"}}, - {"user": {"login": "random-oss-dev"}}, # association field absent - ], - ) - def test_should_fail_safe_when_author_association_is_missing( - self, triage_module, item - ): - # Fail-safe: an empty/missing association must never make a PR - # eligible for the destructive close path. Treat as internal (skip). - assert triage_module.is_internal_contributor(item) is True - - @pytest.mark.parametrize( - "login", - ["dependabot[bot]", "greptile-apps[bot]", "dependabot", "github-actions"], - ) - def test_should_skip_bot_accounts_regardless_of_association( - self, triage_module, login - ): - item = {"author_association": "NONE", "user": {"login": login}} - assert triage_module.is_internal_contributor(item) is True - - -class TestHasLinkedIssue: - @pytest.mark.parametrize( - "body", - [ - "Fixes #1234", - "closes #1", - "Resolves #99", - "fix #42 — this addresses the regression", - "Closes https://github.com/BerriAI/litellm/issues/27000", - "Resolved https://github.com/BerriAI/litellm/issues/27001", - ], - ) - def test_should_detect_common_link_phrases(self, triage_module, body): - assert triage_module.has_linked_issue(body) is True - - @pytest.mark.parametrize( - "body", - [ - "", - "Some change", - # Casual mentions must NOT auto-pass — they should fall through to - # the LLM judge so the stricter "not a passing mention" rule fires. - "See #1234", - "see #1234 for context", - "ref #1234", - "Refs https://github.com/BerriAI/litellm/issues/27000", - "this addresses #1234", - ], - ) - def test_should_not_auto_pass_casual_mentions(self, triage_module, body): - assert triage_module.has_linked_issue(body) is False - - def test_should_not_detect_when_only_html_comment_template(self, triage_module): - body = "" - assert triage_module.has_linked_issue(body) is False - - -class TestStripHtmlComments: - def test_should_remove_single_line_comments(self, triage_module): - text = "before after" - assert "placeholder" not in triage_module.strip_html_comments(text) - - def test_should_remove_multiline_comments(self, triage_module): - text = "kept\n\nkept2" - cleaned = triage_module.strip_html_comments(text) - assert "Fixes #1" not in cleaned - assert "kept" in cleaned and "kept2" in cleaned - - def test_should_handle_none(self, triage_module): - assert triage_module.strip_html_comments(None) == "" - - -class TestCloseCommentText: - """Pin the user-facing language in close comments so changes are intentional.""" - - def test_pr_close_comment_should_recommend_new_pr_primarily(self, triage_module): - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"} - ) - # Primary path: open a new PR (because OSS authors can't reopen a - # bot-closed PR). Secondary path: `@agent-shin reconsider`. - assert "Open a new PR" in body - assert "@agent-shin reconsider" in body - # Old advice that no longer works for OSS contributors must NOT - # appear (they can't reopen a PR closed by a bot/maintainer). - assert "Reopen the PR" not in body - - def test_reopen_comment_should_carry_reconsider_marker(self, triage_module): - # The marker is what the rate-limit guard greps for to detect a - # prior reconsider verdict on the same PR. If the marker ever - # gets dropped from this comment, the cooldown silently breaks - # and a contributor can spam `@agent-shin reconsider` to burn - # LLM budget. - body = triage_module.format_reopen_comment("pr") - assert triage_module.RECONSIDER_COMMENT_MARKER in body - - def test_still_failing_comment_should_carry_reconsider_marker(self, triage_module): - body = triage_module.format_reconsider_still_failing_comment( - "pr", - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"}, - ) - assert triage_module.RECONSIDER_COMMENT_MARKER in body - - def test_pr_close_comment_should_not_promise_automatic_reopen_on_open( - self, triage_module - ): - # The previous comment said "I'll re-evaluate automatically" — that - # only worked because the author could reopen, which they often - # can't. The new wording must point them at the comment trigger or - # a new PR instead. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "I'll re-evaluate automatically" not in body - - def test_issue_close_comment_should_use_reconsider_trigger(self, triage_module): - # OSS authors have read access, which only lets them reopen issues - # they closed themselves; they CANNOT reopen an issue a maintainer or - # bot closed. So the recovery path is `@agent-shin reconsider` (the - # bot reopens), exactly like the PR path. If this regresses to "reopen - # it yourself", contributors hit a dead end on bot-closed issues. - body = triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": ["repro"], "explanation": "thin"} - ) - assert "@agent-shin reconsider" in body - - def test_pr_close_comment_should_link_blog_explainer(self, triage_module): - # The blog post is the canonical public explanation of what the bot - # checks and why. Every action-required bot comment must link to it - # so contributors landing on a bot-closed PR can self-serve context - # without pinging a maintainer. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "https://docs.litellm.ai/blog/agent-shin-triage" in body - - def test_issue_close_comment_should_link_blog_explainer(self, triage_module): - body = triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "https://docs.litellm.ai/blog/agent-shin-triage" in body - - def test_pr_close_comment_should_flag_mocked_tests_as_insufficient_proof( - self, triage_module - ): - # The PR rubric was tightened to require end-to-end QA proof and - # explicitly exclude mocked-dependency unit tests. The user-facing - # close comment must say so — otherwise contributors will keep - # re-submitting "pytest passed (mocks)" runs and getting closed - # again with no explanation of why. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "end-to-end qa proof" in body.lower() - assert "mock" in body.lower() - - def test_issue_recovery_comments_should_name_feature_dead_end_evidence( - self, triage_module - ): - # The feature-request pass bar demands end-to-end evidence of the - # dead-end, so the close and grace-warning recovery bullets must ask - # for it too — otherwise a requester follows those exact instructions - # (description + use case only) and fails `reconsider` again with no - # hint of what else was needed. - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - for body in ( - triage_module.format_issue_close_comment(verdict), - triage_module.format_grace_warning_issue_comment(verdict), - ): - normalized = " ".join(body.split()) - assert "end-to-end evidence of the dead-end" in normalized - assert "showing where the flow stops today" in normalized - - def test_all_agent_shin_comments_should_use_bullet_train_emoji(self, triage_module): - # The bullet train (🚅) is Agent Shin's symbol, matching the LiteLLM - # logo; the previous wave (👋) was generic and didn't match the bot's - # identity. Every action-required comment the bot can post must use the - # bullet train so the contributor recognizes who's writing without - # reading the signoff. - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - comments = { - "pr_close": triage_module.format_pr_close_comment(verdict), - "issue_close": triage_module.format_issue_close_comment(verdict), - "pr_grace": triage_module.format_grace_warning_pr_comment(verdict), - "issue_grace": triage_module.format_grace_warning_issue_comment(verdict), - "within_grace": triage_module.format_within_grace_comment( - [], "", grace_days=1 - ), - } - for name, body in comments.items(): - assert "🚅" in body, f"{name} comment is missing the bullet train emoji" - assert "👋" not in body, f"{name} comment still uses the old wave emoji" - - def test_pr_close_comment_should_show_what_pr_got_right(self, triage_module): - # The user explicitly asked for a "things you got right" section so - # the comment doesn't read as pure rejection. When the judge confirms - # a field is present (e.g. linked_issue), the bullet for it MUST - # appear in the close comment. - body = triage_module.format_pr_close_comment( - { - "verdict": "fail", - "linked_issue": True, - "has_problem_description": True, - "has_expected_vs_actual": False, - "has_qa_proof": False, - "missing": ["QA proof"], - "explanation": "no proof", - } - ) - assert "What you got right" in body - # The two present fields surface as ✅ bullets; the two absent - # fields do not get a ✅ bullet (the QA-proof rubric block still - # mentions the concept, but only the affirmed fields get checkmarks). - assert "- ✅ Linked a related GitHub issue" in body - assert "- ✅ Clear problem description" in body - assert "- ✅ Expected vs. actual behavior" not in body - assert "- ✅ End-to-end QA proof" not in body - - def test_pr_close_comment_should_omit_present_section_when_nothing_present( - self, triage_module - ): - # If the judge says nothing is present (every flag False), the - # "what you got right" block is skipped entirely — better to omit - # than to render "What you got right: (nothing)". - body = triage_module.format_pr_close_comment( - { - "verdict": "fail", - "linked_issue": False, - "has_problem_description": False, - "has_expected_vs_actual": False, - "has_qa_proof": False, - "missing": [], - "explanation": "", - } - ) - assert "What you got right" not in body - - def test_issue_close_comment_should_show_what_issue_got_right(self, triage_module): - # `has_expected_vs_actual` is present, the end-to-end bug evidence is - # not: the "what you got right" block must surface the former and omit - # the latter (no "✅ (nothing)"-style noise for absent items). - body = triage_module.format_issue_close_comment( - { - "verdict": "fail", - "kind": "bug", - "has_repro": False, - "has_expected_vs_actual": True, - "missing": ["end-to-end evidence of the bug"], - "explanation": "no repro shown", - } - ) - assert "What you got right" in body - assert "Expected vs. actual behavior" in body - assert "- ✅ End-to-end evidence of the bug" not in body - - def test_issue_close_comment_should_credit_feature_dead_end_evidence( - self, triage_module - ): - # A feature requester who pasted their dead-end run but skipped the - # motivation must see the evidence credited and only the motivation - # listed as a gap — without a dedicated verdict field the praise - # block could never acknowledge the work they did do. - body = triage_module.format_issue_close_comment( - { - "verdict": "fail", - "kind": "feature", - "has_motivation_example": False, - "has_dead_end_evidence": True, - "missing": ["motivation / use case"], - "explanation": "no use case given", - } - ) - assert "What you got right" in body - assert "- ✅ End-to-end evidence of the dead-end" in body - assert "- ✅ Motivation and concrete example" not in body - - def test_close_comments_should_use_softer_park_for_later_framing( - self, triage_module - ): - # User feedback: the messaging shouldn't feel like punishment. The - # comment must explicitly frame close as a "park this for later," not - # a rejection, and ground that in the queue-hygiene reason. - for body in ( - triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ): - assert "park this for later" in body - assert ( - "not a rejection" in body - or "isn't a rejection" in body - or ("isn't us saying" in body) - ) - - def test_only_close_comments_carry_the_agent_shin_close_marker(self, triage_module): - # The reconsider reopen guard keys off AGENT_SHIN_CLOSE_MARKER to tell - # an Agent Shin close from a same-identity close by another workflow. - # That only works if the marker is stamped on the close comments and - # NOT on the grace warnings (which don't close anything). - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - marker = triage_module.AGENT_SHIN_CLOSE_MARKER - assert marker in triage_module.format_pr_close_comment(verdict) - assert marker in triage_module.format_issue_close_comment(verdict) - assert marker not in triage_module.format_grace_warning_pr_comment(verdict) - assert marker not in triage_module.format_grace_warning_issue_comment(verdict) - - -class TestWasClosedByAgentShin: - """Bot-closed guard: only Agent Shin's own closures are reopen candidates.""" - - @staticmethod - def _stub_close_event( - triage_module, - monkeypatch, - *, - actor: str | None, - closed_at: object = "now", - ): - """Stub the most recent `closed` event used by the guard. - - `actor` is the login that closed the item. `closed_at` defaults - to "now" so the marker comment (stubbed at 42s ago) reads as - recent enough relative to the close; tests can pass a concrete - ``datetime`` to simulate older closes (e.g. the stale-marker - regression scenario). - """ - import datetime as real_dt - - if closed_at == "now": - closed_at = real_dt.datetime.now(real_dt.timezone.utc) - monkeypatch.setattr( - triage_module, - "fetch_last_close_event", - lambda repo, n: (actor, closed_at), - ) - - @staticmethod - def _stub_close_marker_present( - triage_module, monkeypatch, *, present: bool, age_seconds: float = 42.0 - ): - """Stub the Agent Shin close-comment marker lookup. - - `was_closed_by_agent_shin` requires the closing actor AND a - recent Agent Shin close comment; these tests pin the latter so - they exercise the actor half in isolation. - """ - monkeypatch.setattr( - triage_module, - "seconds_since_last_agent_shin_close", - lambda *a, **kw: age_seconds if present else None, - ) - - def test_should_return_true_when_bot_closed_and_close_comment_present( - self, triage_module, monkeypatch - ): - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is True - - def test_should_return_false_when_bot_closed_but_no_agent_shin_comment( - self, triage_module, monkeypatch - ): - # The `github-actions[bot]` identity is shared across workflows. A - # stale/duplicate sweep closing under that identity must NOT let - # @agent-shin reconsider reopen the item: without an Agent Shin close - # comment the guard fails closed. - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - self._stub_close_marker_present(triage_module, monkeypatch, present=False) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_return_false_when_last_close_actor_is_maintainer( - self, triage_module, monkeypatch - ): - # A maintainer closed it (e.g. duplicate, security, design). The - # bot must refuse to reopen on @agent-shin reconsider even if an - # earlier Agent Shin close comment is still on the thread. - self._stub_close_event(triage_module, monkeypatch, actor="krrishdholakia") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_fail_closed_when_no_close_event(self, triage_module, monkeypatch): - # If the events API returns nothing (network blip, repo permission - # quirk), the guard must fail-closed: refuse to reopen rather than - # assume the bot did it. - self._stub_close_event(triage_module, monkeypatch, actor=None, closed_at=None) - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_fail_closed_when_close_event_has_no_timestamp( - self, triage_module, monkeypatch - ): - # Without a usable close timestamp the guard cannot prove the - # marker comment belongs to the latest close; fail-closed. - self._stub_close_event( - triage_module, monkeypatch, actor="github-actions[bot]", closed_at=None - ) - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_return_false_when_marker_predates_latest_close( - self, triage_module, monkeypatch - ): - # Regression for the stale-marker bug: Agent Shin closed once - # (marker stamped), reconsider reopened, and a different workflow - # later closed under the same bot identity without stamping the - # marker. The old marker is still on the thread but does NOT - # belong to the latest close, so reconsider must not reopen. - import datetime as real_dt - - now = real_dt.datetime.now(real_dt.timezone.utc) - # Latest close happened a minute ago. - self._stub_close_event( - triage_module, - monkeypatch, - actor="github-actions[bot]", - closed_at=now - real_dt.timedelta(seconds=60), - ) - # The most recent Agent Shin marker is from an hour ago (a prior - # closed/reopened cycle), which is well outside the skew window. - self._stub_close_marker_present( - triage_module, monkeypatch, present=True, age_seconds=3600.0 - ) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_respect_bot_login_override_via_env( - self, triage_module, monkeypatch - ): - # Operators wiring Agent Shin to a PAT (instead of GITHUB_TOKEN) - # can override the expected bot login via env. The guard must - # respect the override so non-default deployments still work. - monkeypatch.setenv("AGENT_SHIN_BOT_LOGIN", "my-bot") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - self._stub_close_event(triage_module, monkeypatch, actor="my-bot") - assert triage_module.was_closed_by_agent_shin("o/r", 1) is True - # Default "github-actions[bot]" should NOT match when env is set. - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - -class TestSecondsSinceLastAgentShinClose: - """Close-provenance lookup: detects the bot's own auto-close marker.""" - - def _make_comment(self, *, login: str, body: str) -> dict: - return { - "user": {"login": login}, - "body": body, - "created_at": "2026-05-18T05:00:00Z", - } - - def test_should_return_none_when_bot_never_closed(self, triage_module, monkeypatch): - # Comments exist, but none is an Agent Shin close — e.g. only a grace - # warning, or a close by another workflow with no Agent Shin comment. - comments = [ - self._make_comment(login="outside-dev", body="any update?"), - self._make_comment( - login="github-actions[bot]", - body=triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is None - - def test_should_detect_bot_close_comment(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body=triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is not None - - def test_should_ignore_non_bot_comment_quoting_marker( - self, triage_module, monkeypatch - ): - # A contributor quoting the hidden marker (GitHub "Quote reply" - # preserves HTML comments) must not be mistaken for a bot close. - comments = [ - self._make_comment( - login="curious-user", - body=f"what is this? {triage_module.AGENT_SHIN_CLOSE_MARKER}", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is None - - -class TestSecondsSinceLastReconsiderVerdict: - """Rate-limit guard: detects the bot's own reconsider verdict marker.""" - - def _make_comment( - self, *, login: str, body: str, created_at: str | None = "2026-05-18T05:00:00Z" - ) -> dict: - comment: dict = {"user": {"login": login}, "body": body} - if created_at is not None: - comment["created_at"] = created_at - return comment - - def test_should_return_none_when_no_bot_reconsider_comments( - self, triage_module, monkeypatch - ): - # An issue with chatter from other users but no bot reconsider - # verdict must not be rate-limited. - comments = [ - self._make_comment(login="outside-dev", body="ping?"), - self._make_comment( - login="github-actions[bot]", body="some other bot message" - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - def test_should_pick_latest_bot_reconsider_marker(self, triage_module, monkeypatch): - # When multiple reconsider verdicts exist, return the AGE of the - # most recent one. Using a frozen reference helps pin the math. - comments = [ - self._make_comment( - login="github-actions[bot]", - body="old verdict " + triage_module.RECONSIDER_COMMENT_MARKER, - created_at="2026-05-18T04:00:00Z", - ), - self._make_comment( - login="github-actions[bot]", - body="newer verdict " + triage_module.RECONSIDER_COMMENT_MARKER, - created_at="2026-05-18T04:55:00Z", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - - # Freeze "now" via a tiny shim on the module's `dt` import. - import datetime as real_dt - - class FrozenDateTime(real_dt.datetime): - @classmethod - def now(cls, tz=None): - return real_dt.datetime(2026, 5, 18, 5, 0, 0, tzinfo=tz) - - frozen_module = type(triage_module.dt)("datetime") - frozen_module.datetime = FrozenDateTime - frozen_module.timezone = real_dt.timezone - monkeypatch.setattr(triage_module, "dt", frozen_module) - - age = triage_module.seconds_since_last_reconsider_verdict("o/r", 1) - # newer verdict is 5 minutes (300 seconds) before "now" - assert age == 300.0 - - def test_should_ignore_non_bot_comments_with_marker( - self, triage_module, monkeypatch - ): - # A user comment that happens to quote the marker (e.g. in - # a "what does this hidden marker do?" question) must NOT count. - # The rate-limit guard only trusts comments authored by the bot. - comments = [ - self._make_comment( - login="curious-user", - body=f"Saw this marker: {triage_module.RECONSIDER_COMMENT_MARKER}", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - def test_should_ignore_bot_comments_without_marker( - self, triage_module, monkeypatch - ): - # The bot posts other things too (Agent Shin close comments, - # CI status, etc.) — only the reconsider-verdict marker should - # arm the cooldown. - comments = [ - self._make_comment( - login="github-actions[bot]", - body="Agent Shin closed this PR (no marker)", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - -class TestParseVerdict: - def test_should_parse_plain_json(self, triage_module): - raw = '{"verdict": "pass", "missing": []}' - assert triage_module.parse_verdict(raw)["verdict"] == "pass" - - def test_should_strip_markdown_fence(self, triage_module): - raw = '```json\n{"verdict": "fail", "missing": ["foo"]}\n```' - result = triage_module.parse_verdict(raw) - assert result["verdict"] == "fail" - assert result["missing"] == ["foo"] - - def test_should_extract_embedded_json_from_prose(self, triage_module): - raw = 'Here you go: {"verdict": "pass", "missing": []}\nThanks.' - assert triage_module.parse_verdict(raw)["verdict"] == "pass" - - def test_should_raise_for_unparseable_text(self, triage_module): - with pytest.raises(ValueError, match='could not extract JSON from LLM response: not even close to'): - triage_module.parse_verdict("not even close to json") - - def test_should_raise_for_empty(self, triage_module): - with pytest.raises(ValueError, match='empty LLM response'): - triage_module.parse_verdict("") - - -class TestBuildPrompts: - def test_should_include_pr_title_and_body(self, triage_module): - prompt = triage_module.build_pr_prompt( - title="Add foo", body=" Real body" - ) - assert "Add foo" in prompt - assert "Real body" in prompt - assert "comment" not in prompt # HTML comments are stripped - - def test_should_show_empty_marker_for_empty_pr_body(self, triage_module): - prompt = triage_module.build_pr_prompt(title="t", body="") - assert "(empty)" in prompt - - def test_should_include_issue_title_and_body(self, triage_module): - prompt = triage_module.build_issue_prompt(title="Bug", body="repro here") - assert "Bug" in prompt - assert "repro here" in prompt - - def test_issue_bug_rubric_requires_end_to_end_evidence_and_drops_pass_bias( - self, triage_module - ): - # The bug bar was tightened: a report needs the "before" half shown - # end-to-end (video / screenshot / real command output), prose-only - # repro steps no longer pass, and the old "bias toward PASS" leniency - # is gone. If any of these regress, the judge silently goes soft on - # undemonstrated bug reports again. - prompt = triage_module.build_issue_prompt(title="t", body="x") - normalized = " ".join(prompt.split()) - assert "Bias toward PASS when the issue has structure" not in normalized - assert "END-TO-END EVIDENCE OF THE BUG" in normalized - assert "Do not bias toward PASS" in normalized - # The three accepted forms of the "before" demonstration must be named. - assert "screen recording / video" in normalized - assert "screenshot of the bug" in normalized - assert "mocked or stubbed" in normalized - # Prose-only steps are explicitly insufficient now. - assert "steps to reproduce" in normalized - # An unedited issue-form scaffold must not read as evidence: the proof - # field ships with visible headings, so the judge has to be told that - # bare headings with nothing under them count as absent. - assert "unfilled template scaffold" in normalized - assert "counts as absent, not as evidence" in normalized - - def test_issue_feature_rubric_requires_evidence_of_the_dead_end( - self, triage_module - ): - # The feature form asks the requester to walk the ideal flow against a - # live proxy and paste output up to the step that dead-ends, so the - # judge has to demand that evidence, and must not accept an unedited - # scaffold of bare headings as if it were a real attempt. - prompt = triage_module.build_issue_prompt(title="t", body="x") - normalized = " ".join(prompt.split()) - assert "END-TO-END EVIDENCE OF THE DEAD-END" in normalized - assert "showing the point where the flow stops today" in normalized - assert "unfilled template scaffold" in normalized - # The evidence has its own verdict field so feature requesters who - # provided it get credited in "What you got right", exactly like - # `has_repro` credits bug evidence. - assert "`has_dead_end_evidence=true` only when this is present" in normalized - assert '"has_dead_end_evidence": boolean' in normalized - - def test_should_not_crash_when_pr_body_contains_curly_braces(self, triage_module): - """User-supplied content with `{` / `}` must NOT be re-parsed by - `str.format()`. `format` only scans the template literal for - replacement fields; values being substituted in are inserted as - plain strings, so a body like `{"foo": "bar"}` or `{unmatched` - cannot blow up the script. Pinning this here so a future - "improvement" to the templating doesn't reintroduce a crash on - every PR that quotes JSON. - """ - for body in ( - 'Here is some JSON: {"foo": "bar", "n": 1}', - "Half a brace { left dangling, and a stray }", - "Format-spec-looking thing: {0}, {name:>10}, {!r}", - "Nested {a: {b: c}} braces", - ): - pr_prompt = triage_module.build_pr_prompt(title="t", body=body) - issue_prompt = triage_module.build_issue_prompt(title="t", body=body) - assert body in pr_prompt - assert body in issue_prompt - - def test_should_not_crash_when_pr_title_contains_curly_braces(self, triage_module): - title = "Fix bug in {0:>10} format-spec handling" - pr_prompt = triage_module.build_pr_prompt(title=title, body="x") - issue_prompt = triage_module.build_issue_prompt(title=title, body="x") - assert title in pr_prompt - assert title in issue_prompt - - def test_should_preserve_template_indentation_with_multiline_body( - self, triage_module - ): - """`textwrap.dedent` runs on the static template *before* user - content is interpolated, so a multi-line body (whose 2nd+ lines - start at column 0) cannot defeat the common-indent computation - and leave 8-space indentation on every template line. Pin the - dedented shape so the rendered prompt stays consistent for the - LLM judge. - """ - body = "first line\nsecond line at column 0\nthird line at column 0" - for builder in ( - triage_module.build_pr_prompt, - triage_module.build_issue_prompt, - ): - prompt = builder(title="t", body=body) - # Template lines should NOT carry the 8 leading spaces from - # the source-file indentation of the triple-quoted string. - assert " You are " not in prompt - assert 'You are "Agent Shin"' in prompt - assert body in prompt - - -class TestMainModelDefault: - """`--model` falls back to DEFAULT_MODEL even when TRIAGE_MODEL is empty.""" - - def _stub_triage(self, triage_module, monkeypatch): - captured: dict = {} - - def fake_triage(**kwargs): - captured.update(kwargs) - return { - "kind": kwargs["kind"], - "number": kwargs["number"], - "title": "", - "author": "x", - "author_association": "NONE", - "state": "open", - "action": "skip-no-llm-key", - } - - monkeypatch.setattr(triage_module, "triage", fake_triage) - return captured - - def test_should_fall_back_to_default_when_triage_model_env_empty( - self, triage_module, monkeypatch - ): - captured = self._stub_triage(triage_module, monkeypatch) - monkeypatch.setenv("TRIAGE_MODEL", "") - monkeypatch.setattr( - sys, - "argv", - ["triage_with_llm.py", "--repo", "o/r", "--pr", "1"], - ) - rc = triage_module.main() - assert rc == 0 - assert captured["model"] == triage_module.DEFAULT_MODEL - - def test_should_respect_explicit_triage_model_env(self, triage_module, monkeypatch): - captured = self._stub_triage(triage_module, monkeypatch) - monkeypatch.setenv("TRIAGE_MODEL", "gpt-4o-mini") - monkeypatch.setattr( - sys, - "argv", - ["triage_with_llm.py", "--repo", "o/r", "--pr", "1"], - ) - rc = triage_module.main() - assert rc == 0 - assert captured["model"] == "gpt-4o-mini" - - -class TestCallLlmJudge: - """call_llm_judge sets gpt-5 specific kwargs correctly.""" - - def _stub_openai(self, monkeypatch, captured: dict): - """Install a fake `openai.OpenAI` client into sys.modules. - - The fake client records the kwargs passed to chat.completions.create - and returns a minimal response object whose .choices[0].message.content - is "ok". - """ - import types - - class FakeMessage: - content = '{"verdict": "pass"}' - - class FakeChoice: - message = FakeMessage() - - class FakeResponse: - choices = [FakeChoice()] - - class FakeCompletions: - def create(self, **kwargs): - captured.update(kwargs) - return FakeResponse() - - class FakeChat: - completions = FakeCompletions() - - class FakeClient: - def __init__(self, api_key, base_url=None): - captured["__client_kwargs__"] = { - "api_key": api_key, - "base_url": base_url, - } - self.chat = FakeChat() - - fake_module = types.ModuleType("openai") - fake_module.OpenAI = FakeClient - monkeypatch.setitem(sys.modules, "openai", fake_module) - - def test_should_set_reasoning_effort_none_for_gpt5_family( - self, triage_module, monkeypatch - ): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model="gpt-5.4-mini", api_key="sk-test", base_url=None - ) - assert captured["model"] == "gpt-5.4-mini" - assert captured["temperature"] == 0 - assert captured["extra_body"] == {"reasoning_effort": "none"} - - def test_should_set_reasoning_effort_for_capitalized_or_dated_gpt5( - self, triage_module, monkeypatch - ): - for model in ("GPT-5.4-mini", "gpt-5.4-mini-2026-03-17", "gpt-5"): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model=model, api_key="sk-test", base_url=None - ) - assert captured["extra_body"] == {"reasoning_effort": "none"}, model - - def test_should_omit_reasoning_effort_for_non_gpt5( - self, triage_module, monkeypatch - ): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model="gpt-4o-mini", api_key="sk-test", base_url=None - ) - assert "extra_body" not in captured - - def test_should_pass_base_url_when_provided(self, triage_module, monkeypatch): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "p", - model="gpt-5.4-mini", - api_key="sk-test", - base_url="https://proxy.example.com/v1", - ) - assert ( - captured["__client_kwargs__"]["base_url"] == "https://proxy.example.com/v1" - ) - - -class TestTriageOrchestration: - """End-to-end-ish tests that mock both gh fetchers and the LLM.""" - - def _make_pr(self, **overrides): - base = { - "number": 1, - "title": "PR title", - "body": "PR body", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - base.update(overrides) - return base - - def test_should_skip_internal_author(self, triage_module, monkeypatch): - pr = self._make_pr( - author_association="MEMBER", user={"login": "krrishdholakia"} - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - - def boom(*a, **kw): - pytest.fail("LLM should not be called for internal authors") - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=boom, - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_should_skip_closed_pr(self, triage_module, monkeypatch): - pr = self._make_pr(state="closed") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("should not run on closed PRs"), - ) - assert result["action"] == "skip-not-open" - - def test_should_short_circuit_on_linked_issue(self, triage_module, monkeypatch): - pr = self._make_pr(body="Fixes #1234\n\nFoo bar") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM should not be called"), - ) - assert result["action"] == "pass-linked-issue" - assert result["verdict"]["verdict"] == "pass" - - def test_should_not_short_circuit_on_casual_mention( - self, triage_module, monkeypatch - ): - # "See #1234" is a passing mention, not a closing keyword. The LLM - # must get a chance to apply the stricter rubric. With no prior - # grace warning, the first failing verdict triggers the warning - # path (`would-warn-grace` in dry-run). - pr = self._make_pr(body="See #1234 for context. No QA proof here.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - called = {"judge": False} - - def judge(prompt): - called["judge"] = True - return json.dumps( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin."} - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=judge, - ) - assert called["judge"] is True - assert result["action"] == "would-warn-grace" - - def test_should_return_pass_llm_when_judge_passes(self, triage_module, monkeypatch): - pr = self._make_pr(body="Long body, no linked issue.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - captured = {} - - def judge(prompt): - captured["prompt"] = prompt - return json.dumps({"verdict": "pass", "missing": [], "explanation": "ok"}) - - result = triage_module.triage( - repo="o/r", kind="pr", number=1, close=True, model="m", judge=judge - ) - assert result["action"] == "pass-llm" - assert "Long body" in captured["prompt"] - - def test_should_return_would_close_in_dry_run_after_grace_aged_out( - self, triage_module, monkeypatch - ): - # When the grace warning has already aged out (>= GRACE_PERIOD_SECONDS) - # AND the rubric still fails, the dry-run preview returns - # `would-close` so a step-summary writer can render the close - # comment without touching GitHub state. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_aged_out(triage_module, monkeypatch) - - def fake_post(*a, **kw): - pytest.fail("should not post comments in dry-run") - - def fake_close(*a, **kw): - pytest.fail("should not close in dry-run") - - monkeypatch.setattr(triage_module, "post_comment", fake_post) - monkeypatch.setattr(triage_module, "close_pr", fake_close) - - verdict = { - "verdict": "fail", - "missing": ["problem description", "QA proof"], - "explanation": "Body is one sentence.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "would-close" - assert result["verdict"]["missing"] == ["problem description", "QA proof"] - - def test_should_post_comment_and_close_after_grace_window( - self, triage_module, monkeypatch - ): - # The "real close" path: --close passed AND the grace warning has - # aged out AND the rubric still fails. The bot posts the close - # comment and closes the PR. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_aged_out(triage_module, monkeypatch) - posted = {} - closed = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"repo": repo, "n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda repo, n: closed.update({"repo": repo, "n": n}), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "closed" - assert posted["n"] == 42 and closed["n"] == 42 - assert "Agent Shin" in posted["body"] - assert "QA proof" in posted["body"] - - def test_should_skip_on_llm_error_in_close_mode(self, triage_module, monkeypatch): - pr = self._make_pr(body="something.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment on LLM error"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on LLM error"), - ) - - def broken_judge(prompt): - raise RuntimeError("upstream 500") - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=broken_judge, - ) - assert result["action"] == "skip-llm-error" - assert "upstream 500" in result["error"] - - def test_should_skip_open_pr_in_reconsider_mode(self, triage_module, monkeypatch): - # Reconsider only makes sense on a CLOSED PR — running it on an open - # one is a no-op (the regular triage flow already evaluated it). - pr = self._make_pr(state="open") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: pytest.fail("should not run on open PR in reconsider"), - reconsider=True, - ) - assert result["action"] == "skip-not-closed" - - @staticmethod - def _stub_reconsider_guards(triage_module, monkeypatch): - """Default reconsider-guard stubs: pretend bot closed + no cooldown. - - The new safety guards (`was_closed_by_agent_shin`, - `seconds_since_last_reconsider_verdict`) hit the GitHub API in - production. Tests that exercise the reconsider happy path stub - them to "yes the bot closed it, no recent reconsider comment" - so the test stays focused on its actual assertion. - """ - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: None, - ) - - @staticmethod - def _stub_grace_aged_out(triage_module, monkeypatch): - """Pretend the grace warning has aged out. - - For tests that exercise the post-grace close path. Set the age - to twice the grace window so a future tweak to - `GRACE_PERIOD_SECONDS` doesn't accidentally make the stub fall - back inside the window. - """ - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: triage_module.GRACE_PERIOD_SECONDS * 2, - ) - - @staticmethod - def _stub_grace_no_warning(triage_module, monkeypatch): - """Pretend no grace warning has been posted yet (first detection).""" - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: None, - ) - - def test_should_reopen_on_reconsider_pass(self, triage_module, monkeypatch): - # Reconsider on a closed PR with a passing verdict -> reopen + post a - # friendly "re-evaluated" comment. close=True is the production path - # (the workflow only adds --close when AGENT_SHIN_ENABLED=true). - pr = self._make_pr( - state="closed", body="Updated body with QA proof + screenshots." - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - # close_pr / close_issue MUST NOT fire in reconsider mode. - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on reconsider pass"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok now"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 42 - assert posted["n"] == 42 - assert "reopened" in posted["body"].lower() - - def test_should_dry_run_reconsider_pass_when_close_false( - self, triage_module, monkeypatch - ): - # Reconsider must honor `close=False` (dry-run) just like the - # regular triage flow. A local invocation of - # `python triage_with_llm.py --reconsider --pr N` (no --close) - # must NOT post a comment or reopen the PR — it should return - # `would-reopen` so the operator can preview the outcome. - pr = self._make_pr( - state="closed", body="Updated body with QA proof + screenshots." - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post comment in dry-run reconsider"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen PR in dry-run reconsider"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=False, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok now"} - ), - reconsider=True, - ) - assert result["action"] == "would-reopen" - # The previewed comment body is still returned so a step-summary - # writer can render exactly what would have been posted. - assert "reopened" in result["comment"].lower() - - def test_should_post_still_failing_on_reconsider_fail( - self, triage_module, monkeypatch - ): - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - # Neither reopen nor close should fire when reconsider verdict is fail. - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen on fail"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close again on reconsider fail"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Still no QA proof.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - reconsider=True, - ) - assert result["action"] == "reconsider-still-failing" - assert posted["n"] == 42 - assert "QA proof" in posted["body"] - - def test_should_not_reopen_on_reconsider_with_ambiguous_verdict( - self, triage_module, monkeypatch - ): - # Regression: only an explicit `pass` verdict reopens. Missing, - # empty, or unexpected verdict strings ("failed", "", garbage) - # must fall through to the still-failing branch rather than - # reopen a PR the rubric did not actually clear. - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen on ambiguous verdict"), - ) - - for ambiguous in ("", "failed", "needs-info", "unknown"): - posted.clear() - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p, v=ambiguous: json.dumps( - {"verdict": v, "missing": [], "explanation": "weird"} - ), - reconsider=True, - ) - assert result["action"] == "reconsider-still-failing", ambiguous - assert "body" in posted, ambiguous - - def test_should_dry_run_reconsider_fail_when_close_false( - self, triage_module, monkeypatch - ): - # Mirror dry-run behavior for the FAIL branch — `close=False` - # must NOT post the "still failing" comment. - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail( - "must not post still-failing comment in dry-run" - ), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Still no QA proof.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - reconsider=True, - ) - assert result["action"] == "would-reconsider-still-failing" - assert "QA proof" in result["comment"] - - def test_should_reopen_on_reconsider_with_linked_issue_short_circuit( - self, triage_module, monkeypatch - ): - # The linked-issue short-circuit also has to honor reconsider mode: - # if the contributor edited the body to add `Fixes #1234`, the regex - # path should reopen the PR without calling the LLM. - pr = self._make_pr(state="closed", body="Fixes #1234\n\nAddresses the bug.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=55, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run when linked-issue matches"), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 55 - assert "reopened" in posted["body"].lower() - - def test_should_dry_run_reconsider_with_linked_issue_when_close_false( - self, triage_module, monkeypatch - ): - # Linked-issue short-circuit must ALSO honor dry-run. - pr = self._make_pr(state="closed", body="Fixes #1234\n\nAddresses the bug.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post in dry-run"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen in dry-run"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=55, - close=False, - model="m", - judge=lambda p: pytest.fail("LLM must not run when linked-issue matches"), - reconsider=True, - ) - assert result["action"] == "would-reopen" - - def test_should_skip_internal_in_reconsider_mode(self, triage_module, monkeypatch): - # Internal authors are exempt from triage in both regular and - # reconsider mode — Agent Shin should never reopen one of their PRs - # automatically, in case a maintainer closed it intentionally. - pr = self._make_pr( - state="closed", - author_association="MEMBER", - user={"login": "krrishdholakia"}, - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen for internal author"), - ) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: pytest.fail("LLM must not run for internal author"), - reconsider=True, - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_should_skip_reconsider_when_not_bot_closed( - self, triage_module, monkeypatch - ): - # SECURITY: `@agent-shin reconsider` must NOT reopen a PR/issue - # that a MAINTAINER closed for non-rubric reasons (e.g. duplicate, - # design rejection, security report). Only PRs closed by the bot - # itself should ever be candidates for the reconsider reopen path. - pr = self._make_pr(state="closed", body="something.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: False - ) - # Even though there's no rate-limit conflict, the bot-closed guard - # alone is sufficient to block. The LLM judge must never run on a - # maintainer-closed PR. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: None, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment on maintainer-closed PR"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen maintainer-closed PR"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run before bot-closed guard"), - reconsider=True, - ) - assert result["action"] == "skip-not-bot-closed" - - def test_should_rate_limit_repeated_reconsider_triggers( - self, triage_module, monkeypatch - ): - # COST CONTROL: each `@agent-shin reconsider` event burns CI - # minutes + an OpenAI API call. If the bot already posted a - # reconsider verdict within the cooldown window - # (RECONSIDER_RATE_LIMIT_SECONDS), refuse to run again. This - # bounds the damage from a contributor spamming the trigger. - pr = self._make_pr(state="closed", body="something with new edits.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - # Pretend the bot posted a reconsider verdict 1 second ago. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: 1.0, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment during cooldown"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen during cooldown"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run during cooldown"), - reconsider=True, - ) - assert result["action"] == "skip-rate-limited" - assert result["rate_limit_age_seconds"] == 1.0 - assert ( - result["rate_limit_window_seconds"] - == triage_module.RECONSIDER_RATE_LIMIT_SECONDS - ) - - def test_should_allow_reconsider_after_cooldown_window( - self, triage_module, monkeypatch - ): - # The cooldown is a window, not a one-shot lock — once - # RECONSIDER_RATE_LIMIT_SECONDS has elapsed since the last bot - # verdict, a fresh `@agent-shin reconsider` is allowed through. - pr = self._make_pr(state="closed", body="updated with screenshots now.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - # Last reconsider was 1 hour ago — well outside the 10-min window. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: 3600.0, - ) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 1 - - def test_should_reopen_issue_on_reconsider_pass(self, triage_module, monkeypatch): - issue = { - "number": 7, - "title": "Bug: now with repro", - "body": "## Repro\n```bash\ncurl ...\n```\n\nExpected X, got Y.", - "state": "closed", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - monkeypatch.setattr(triage_module, "fetch_issue", lambda repo, n: issue) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_issue", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="issue", - number=7, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "now reproducible"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 7 - assert "reopened" in posted["body"].lower() - - def test_should_triage_issues_kind(self, triage_module, monkeypatch): - issue = { - "number": 7, - "title": "Bug: X is broken", - "body": "no detail", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - monkeypatch.setattr(triage_module, "fetch_issue", lambda repo, n: issue) - # Grace already aged out -> close path. (Issues use the same - # GRACE_COMMENT_MARKER detection as PRs.) - self._stub_grace_aged_out(triage_module, monkeypatch) - closed = {} - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update(body=body), - ) - monkeypatch.setattr( - triage_module, "close_issue", lambda repo, n: closed.update(n=n) - ) - - verdict = { - "verdict": "fail", - "kind": "bug", - "has_repro": False, - "missing": ["reproduction", "expected vs. actual"], - "explanation": "No repro provided.", - } - result = triage_module.triage( - repo="o/r", - kind="issue", - number=7, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "closed" - assert closed["n"] == 7 - assert "reproduction" in posted["body"] - - # ---- Grace-period flow ------------------------------------------------ - - def test_should_post_grace_warning_on_first_failing_run_in_close_mode( - self, triage_module, monkeypatch - ): - # First low-quality detection -> bot posts a warning comment with - # the GRACE_COMMENT_MARKER. The PR must NOT be closed yet. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on first detection"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "warned-grace" - assert posted["n"] == 42 - # Pin the user-facing language pieces the user explicitly asked for. - assert "2 hours" in posted["body"] - assert "@agent-shin reconsider" in posted["body"] - assert "@greptileai" in posted["body"] - assert "even after the PR is closed" in posted["body"] - assert triage_module.GRACE_COMMENT_MARKER in posted["body"] - - def test_should_skip_close_inside_grace_window(self, triage_module, monkeypatch): - # A warning was posted recently; do nothing on this run regardless - # of close=True. The next run after `GRACE_PERIOD_SECONDS` elapses - # is the one that flips to actual close. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: 60.0, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment during grace window"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close during grace window"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "skip-in-grace-period" - assert result["grace_age_seconds"] == 60.0 - assert result["grace_period_seconds"] == triage_module.GRACE_PERIOD_SECONDS - - def test_should_dry_run_grace_warning_when_close_false( - self, triage_module, monkeypatch - ): - # In dry-run mode the FIRST failing detection returns - # `would-warn-grace` (with the previewed comment body) and never - # touches GitHub state. Lets a local operator preview the - # warning before flipping --close on. - pr = self._make_pr(body="thin") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post in dry-run grace warn"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "thin", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "would-warn-grace" - assert "2 hours" in result["comment"] - - def test_should_warn_grace_for_swiftwinds_not_close_instantly( - self, triage_module, monkeypatch - ): - # Regression: SwiftWinds (the dogfood account) used to be in a - # now-removed `IMMEDIATE_CLOSE_LOGINS` bypass that skipped the grace - # window and closed on first detection. It must follow the SAME - # grace path as every other author: warn first, close only after the - # window elapses. A re-added instant-close bypass would call - # close_pr here and fail the test. - pr = self._make_pr(body="just a sentence.", user={"login": "SwiftWinds"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail( - "SwiftWinds must not close on first detection; it gets the grace window" - ), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=99, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "warned-grace" - assert "2 hours" in posted["body"] - - -class TestGraceWarningCommentText: - """Pin the user-facing promises in the grace warning so a future - refactor can't silently drop them.""" - - def test_pr_grace_warning_should_state_grace_window(self, triage_module): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"} - ) - # The user explicitly asked: "specify in the comment" the grace window. - assert "2 hours" in body - - def test_pr_grace_warning_should_mention_reconsider_during_grace( - self, triage_module - ): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "@agent-shin reconsider" in body - - def test_pr_grace_warning_should_promise_greptileai_works_post_close( - self, triage_module - ): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - # Per user: comment should state @greptileai works even after close. - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_pr_grace_warning_should_carry_grace_marker(self, triage_module): - # The marker is what `seconds_since_last_grace_warning` greps for - # on subsequent runs to detect that a warning has been posted. - # Dropping it would silently break the close-after-grace path. - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert triage_module.GRACE_COMMENT_MARKER in body - - def test_issue_grace_warning_should_carry_grace_marker(self, triage_module): - body = triage_module.format_grace_warning_issue_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert triage_module.GRACE_COMMENT_MARKER in body - assert "2 hours" in body - # OSS authors can't reopen a bot-closed issue, so recovery is - # `@agent-shin reconsider` (the bot reopens), like the PR path. - assert "@agent-shin reconsider" in body - - def test_pr_close_comment_should_promise_greptileai_works_post_close( - self, triage_module - ): - # The standard close comment must ALSO point at @greptileai so - # contributors see the same options whether they read the warning - # or only catch the close comment. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_pr_grace_warning_should_not_prompt_reconsider_during_grace_window( - self, triage_module - ): - # Per user feedback: during the 24h grace window, the contributor - # should just update the PR description. Asking them to also comment - # "@agent-shin reconsider" right away adds a step they don't need — - # the bot re-checks automatically on the next sweep. The reconsider - # trigger is reserved for the post-close recovery path. - # - # We pin this by checking that the grace section explicitly tells - # the contributor they don't need to ping the bot during the grace - # window. The presence of "@agent-shin reconsider" elsewhere in the - # comment (as the post-close path) is fine and required by other - # tests. - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "No need to ping" in body or "no need to ping" in body - - def test_grace_warnings_should_show_what_got_right(self, triage_module): - # The "What you got right" section must appear in the grace warning - # too, not only the close comment — the contributor sees the warning - # first and that's their best chance to know what to keep. - pr_body = triage_module.format_grace_warning_pr_comment( - { - "verdict": "fail", - "linked_issue": True, - "has_problem_description": True, - "has_expected_vs_actual": True, - "has_qa_proof": False, - "missing": ["QA proof"], - "explanation": "thin", - } - ) - assert "What you got right" in pr_body - assert "Linked a related GitHub issue" in pr_body - - issue_body = triage_module.format_grace_warning_issue_comment( - { - "verdict": "fail", - "kind": "feature", - "has_motivation_example": True, - "missing": ["concrete description"], - "explanation": "vague", - } - ) - assert "What you got right" in issue_body - assert "Motivation and concrete example" in issue_body - - def test_grace_warnings_should_use_softer_park_for_later_framing( - self, triage_module - ): - # Same softer-framing pin as the close comment, but for the warning - # — the contributor's first contact with the bot must not read as a - # hard deadline / ultimatum. - for body in ( - triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - triage_module.format_grace_warning_issue_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ): - assert "park this for later" in body - assert ( - "not a rejection" in body - or "isn't a rejection" in body - or ("isn't us saying" in body) - ) - - -class TestSecondsSinceLastGraceWarning: - """Mirror of TestSecondsSinceLastReconsiderVerdict for the new helper. - Both helpers share `_seconds_since_latest_marker_comment` underneath - so the parsing logic is exercised either way; these tests pin the - grace-marker-specific behavior.""" - - def _make_comment( - self, - *, - login: str, - body: str, - created_at: str | None = "2026-05-18T05:00:00Z", - ) -> dict: - comment: dict = {"user": {"login": login}, "body": body} - if created_at is not None: - comment["created_at"] = created_at - return comment - - def test_should_return_none_when_no_grace_marker(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body="Some other bot message", - ), - self._make_comment(login="random-user", body="ping?"), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_grace_warning("o/r", 1) is None - - def test_should_ignore_non_bot_comments_with_marker( - self, triage_module, monkeypatch - ): - # A user who quotes the marker in a question must NOT be treated - # as the bot warning; otherwise the close-after-grace path would - # never fire because the timer keeps resetting. - comments = [ - self._make_comment( - login="random-user", - body=f"What is {triage_module.GRACE_COMMENT_MARKER}?", - ) - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_grace_warning("o/r", 1) is None - - def test_should_pick_latest_grace_marker(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body="old warning " + triage_module.GRACE_COMMENT_MARKER, - created_at="2026-05-18T03:00:00Z", - ), - self._make_comment( - login="github-actions[bot]", - body="newer warning " + triage_module.GRACE_COMMENT_MARKER, - created_at="2026-05-18T04:55:00Z", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - - import datetime as real_dt - - class FrozenDateTime(real_dt.datetime): - @classmethod - def now(cls, tz=None): - return real_dt.datetime(2026, 5, 18, 5, 0, 0, tzinfo=tz) - - frozen_module = type(triage_module.dt)("datetime") - frozen_module.datetime = FrozenDateTime - frozen_module.timezone = real_dt.timezone - monkeypatch.setattr(triage_module, "dt", frozen_module) - - age = triage_module.seconds_since_last_grace_warning("o/r", 1) - # Newer warning is 5 minutes (300s) before "now". - assert age == 300.0 - - -class TestTriageAllowlist: - """The dogfood allowlist gates `triage`: while non-empty it is the sole - author filter (only the named accounts are acted on) and it bypasses the - internal-author exemption for them, so a maintainer can dogfood on their - own org account. Emptying it restores the internal-author skip.""" - - def _make_pr(self, **overrides): - base = { - "number": 1, - "title": "PR title", - "body": "Body with no linked issue and no QA proof.", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - base.update(overrides) - return base - - def test_should_skip_author_not_on_allowlist(self, triage_module, monkeypatch): - pr = self._make_pr(user={"login": "random-oss-dev"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run for non-allowlisted author"), - ) - assert result["action"] == "skip-not-allowlisted" - - def test_should_act_on_allowlisted_internal_author( - self, triage_module, monkeypatch - ): - pr = self._make_pr(author_association="MEMBER", user={"login": "mateo-berri"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok"} - ), - ) - assert result["action"] == "pass-llm" - - def test_empty_allowlist_restores_internal_skip(self, triage_module, monkeypatch): - pr = self._make_pr( - author_association="MEMBER", user={"login": "krrishdholakia"} - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run for internal author"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_allowlist_constant_is_the_two_dogfood_accounts(self, triage_module): - assert triage_module.ALLOWLIST_LOGINS == frozenset( - {"mateo-berri", "swiftwinds"} - ) - for login in triage_module.ALLOWLIST_LOGINS: - assert login == login.lower(), login diff --git a/tests/test_litellm/test_github_triage_workflows.py b/tests/test_litellm/test_github_triage_workflows.py deleted file mode 100644 index f96c9b7e974..00000000000 --- a/tests/test_litellm/test_github_triage_workflows.py +++ /dev/null @@ -1,264 +0,0 @@ -"""Static guardrails for the Agent Shin + Greptile workflow YAML files. - -These workflows can post comments and close PRs/issues on -BerriAI/litellm, so the gating logic that decides "is this a real -close-on-fail run?" must fail-safe on any unexpected input. The risk -is mostly maintenance: someone edits the bash gate, drops a quote, -inverts a comparison, or uses `!= "false"` (which treats "True", -"yes", "1", and typos as enabling closure) and the regression isn't -caught until a real OSS contributor's PR gets auto-closed. - -The tests below pin a set of invariants. The first two apply to every -workflow that gates a destructive `--close`: - - 1. The gate uses the fail-safe `= "true"` comparison — not `!= "false"`, - not `!= ""`. Only the literal string "true" should ever enable - closure. - 2. The gate also requires `AGENT_SHIN_ENABLED = "true"` (or the - scheduled-job equivalent) — disabling the variable must always - force dry-run. - -A third invariant covers every workflow that installs the OpenAI client. -These run with a write-scoped `GITHUB_TOKEN`, so a compromised package -release would execute in that context; the install must therefore come -from the hash-pinned `.github/scripts/triage-requirements.txt` via -`pip --require-hashes`, never a floating `pip install openai>=...`. - -Static parsing of the YAML + bash text is the right level of test here: -the gating logic lives in a `run:` block, not in a Python module we can -import, and end-to-end testing a GitHub Actions workflow from CI is -infeasible. A YAML-level guardrail is exactly what would have caught -the original `!= "false"` regression at PR time. -""" - -from __future__ import annotations - -from pathlib import Path - -import pytest -import yaml - -REPO_ROOT = Path(__file__).resolve().parents[2] -WORKFLOWS_DIR = REPO_ROOT / ".github" / "workflows" - -# Map of workflow file -> the env var name that drives the destructive -# gate inside that workflow's `run:` block. Keeping this table explicit -# (rather than scraping every workflow file) means a new workflow file -# that bypasses the dry-run gating doesn't silently slip past this test. -DESTRUCTIVE_GATE_ENV: dict[str, str] = { - "close_low_quality_prs.yml": "CLOSE_FLAG", - # The reconsider workflow has no per-run "really do it?" knob — its - # only kill switch is `AGENT_SHIN_ENABLED`, which already serves as - # both the destructive gate and the global enablement gate. - "triage_reconsider.yml": "AGENT_SHIN_ENABLED", -} - - -# Privileged workflows that install the OpenAI client. They run with a -# write-scoped GITHUB_TOKEN, so the install must be hash-pinned: a poisoned -# release would otherwise execute in that context. A new workflow that -# installs the client must be added here and use the same pinned file. -LLM_CLIENT_INSTALLER_WORKFLOWS = ( - "triage_reconsider.yml", -) - -PINNED_INSTALL = "--require-hashes -r .github/scripts/triage-requirements.txt" -REQUIREMENTS_FILE = REPO_ROOT / ".github" / "scripts" / "triage-requirements.txt" - - -def _load_workflow(name: str) -> dict: - return yaml.safe_load((WORKFLOWS_DIR / name).read_text()) - - -def _all_run_blocks(workflow: dict) -> list[str]: - """Return every `run:` step's command text, joined.""" - commands: list[str] = [] - jobs = workflow.get("jobs") or {} - for job in jobs.values(): - for step in job.get("steps", []) or []: - if not isinstance(step, dict): - continue - run = step.get("run") - if isinstance(run, str): - commands.append(run) - return commands - - -@pytest.mark.parametrize("workflow_file,env_var", sorted(DESTRUCTIVE_GATE_ENV.items())) -def test_should_use_failsafe_equals_true_comparison(workflow_file: str, env_var: str) -> None: - """The destructive `--close` gate must use `= "true"` (fail-safe), not - `!= "false"` (which would treat "True", "yes", "1", or any typo as - enabling closure). - - Both bare `${ENV_VAR}` and `${ENV_VAR:-false}` (with a default) are - accepted forms — what matters is the comparison operator. The - Greptile closer relies on an outer `AGENT_SHIN_ENABLED` gate so it - can use the bare form; the Agent Shin workflows include `:-false` - for defense in depth. Either is fine. - """ - workflow = _load_workflow(workflow_file) - text = "\n".join(_all_run_blocks(workflow)) - assert env_var in text, ( - f"{workflow_file} no longer references {env_var}; was the gating env var renamed without updating this test?" - ) - accepted_patterns = ( - f'"${{{env_var}}}" = "true"', - f'"${{{env_var}:-false}}" = "true"', - ) - assert any(p in text for p in accepted_patterns), ( - f"{workflow_file} must gate the destructive --close flag on the " - f'EXACT string "true" (one of: {accepted_patterns!r}). Mirror ' - 'the Greptile closer pattern; do NOT use `!= "false"` which ' - 'fail-opens on unknown values like "True", "yes", "1", or typos.' - ) - forbidden_patterns = ( - f'"${{{env_var}}}" != "false"', - f'"${{{env_var}:-false}}" != "false"', - f'"${{{env_var}:-true}}" != "false"', - ) - for forbidden in forbidden_patterns: - assert forbidden not in text, ( - f"{workflow_file} uses the fail-open pattern {forbidden!r}. " - 'Switch to `= "true"` so unknown values stay dry-run.' - ) - - -@pytest.mark.parametrize("workflow_file", sorted(DESTRUCTIVE_GATE_ENV)) -def test_should_require_agent_shin_enabled_for_close(workflow_file: str) -> None: - """Every destructive gate must also gate on the global enablement - variable, so flipping `AGENT_SHIN_ENABLED` off is a kill switch - regardless of any per-run input. - - Two patterns are equally fine: - - Positive: `[ "${AGENT_SHIN_ENABLED:-false}" = "true" ]` to enter - the close branch (Agent Shin workflows). - - Negative: `[ "${AGENT_SHIN_ENABLED:-false}" != "true" ]` then - bail out / force dry-run (Greptile closer). - - What matters is that the comparison value is the literal "true"; - `!= "false"` or `= "1"` etc. would not be a true kill switch. - """ - workflow = _load_workflow(workflow_file) - text = "\n".join(_all_run_blocks(workflow)) - accepted_patterns = ( - '"${AGENT_SHIN_ENABLED:-false}" = "true"', - '"${AGENT_SHIN_ENABLED:-false}" != "true"', - ) - assert any(p in text for p in accepted_patterns), ( - f"{workflow_file} must gate destructive actions on " - '`AGENT_SHIN_ENABLED = "true"` (or the inverted `!= "true"` ' - "guard that forces dry-run). Without this, an unset repo " - "variable would not be treated as a kill switch." - ) - - -@pytest.mark.parametrize("workflow_file", LLM_CLIENT_INSTALLER_WORKFLOWS) -def test_llm_client_install_is_hash_pinned(workflow_file: str) -> None: - """Every privileged workflow installs the OpenAI client from the - hash-pinned requirements file, never by floating version. - - A bare `pip install "openai>=1.40.0"` resolves to whatever PyPI serves - at run time and executes during install/import while a write-scoped - `GITHUB_TOKEN` is in scope, so a compromised release runs in a - privileged context. This test fails if that floating form comes back or - if the `--require-hashes` install is loosened. - """ - blocks = _all_run_blocks(_load_workflow(workflow_file)) - assert PINNED_INSTALL in "\n".join(blocks), ( - f"{workflow_file} must install the client via `pip install " - f"{PINNED_INSTALL}`; a floating install runs unverified code with a " - "write-scoped token." - ) - offenders = [b for b in blocks if "pip install" in b and "openai" in b] - assert not offenders, ( - f"{workflow_file} installs openai by name ({offenders!r}); pin it " - "through the hash-locked requirements file so the version and " - "checksum are fixed." - ) - - -def test_triage_requirements_are_fully_hash_pinned() -> None: - """The shared requirements file pins every package to an exact version - with a sha256 hash, which is what `pip --require-hashes` enforces at - install time. A loosened pin or a missing hash here would silently widen - the supply-chain surface for all the installer workflows. - """ - assert REQUIREMENTS_FILE.exists(), ( - f"the hash-pinned requirements file the triage workflows install from is missing at {REQUIREMENTS_FILE}" - ) - joined = REQUIREMENTS_FILE.read_text().replace("\\\n", " ") - entries = [line.strip() for line in joined.splitlines() if line.strip() and not line.strip().startswith("#")] - assert any(e.split()[0].startswith("openai==") for e in entries), ( - "openai must be pinned to an exact version in the triage requirements" - ) - for entry in entries: - spec = entry.split()[0] - assert "==" in spec, ( - f"requirement {spec!r} is not pinned to an exact version; " - "--require-hashes needs every package pinned with ==" - ) - assert "--hash=sha256:" in entry, ( - f"requirement {spec!r} has no sha256 hash; every pin must carry " - "checksums so --require-hashes can verify the download" - ) - - -def _reconsider_steps() -> list[dict]: - workflow = _load_workflow("triage_reconsider.yml") - return workflow["jobs"]["reconsider"]["steps"] - - -def _index_of_run_step(steps: list[dict], needle: str) -> int: - for i, step in enumerate(steps): - run = step.get("run") - if isinstance(run, str) and needle in run: - return i - raise AssertionError(f"no run step contains {needle!r}") - - -def _reaction_steps(steps: list[dict], content: str) -> list[tuple[int, dict]]: - return [ - (i, s) - for i, s in enumerate(steps) - if isinstance(s.get("run"), str) and f"content={content}" in s["run"] and "/reactions" in s["run"] - ] - - -class TestReconsiderReactions: - """The reconsider workflow acknowledges the triggering comment with a 👀 - reaction the moment it accepts the trigger, and a 👍 once the run finishes, - so the contributor gets feedback immediately instead of waiting on a cron. - - Both reactions are gated on `AGENT_SHIN_ENABLED == 'true'` so a dry-run - leaves no visible trace, and both target the comment that fired the event - (`github.event.comment.id`). The ordering (👀 before the triage run, 👍 - after) is the whole point — these tests fail if a refactor reorders the - steps, drops a reaction, or stops gating them. - """ - - def test_eyes_reaction_is_posted_before_the_triage_run(self) -> None: - steps = _reconsider_steps() - run_idx = _index_of_run_step(steps, "triage_with_llm.py") - eyes = _reaction_steps(steps, "eyes") - assert len(eyes) == 1, "expected exactly one 👀 (eyes) reaction step" - idx, step = eyes[0] - assert idx < run_idx, "👀 must be posted BEFORE the slow triage run, not after" - assert "github.event.comment.id" in (step.get("env") or {}).get("COMMENT_ID", ""), ( - "👀 must react to the comment that triggered the workflow" - ) - assert "${COMMENT_ID}" in step["run"], "👀 must react to the triggering comment, not a hardcoded id" - assert "vars.AGENT_SHIN_ENABLED == 'true'" in step["if"], ( - "👀 must be gated on AGENT_SHIN_ENABLED so dry-run stays inert" - ) - - def test_thumbs_up_reaction_is_posted_after_a_successful_run(self) -> None: - steps = _reconsider_steps() - run_idx = _index_of_run_step(steps, "triage_with_llm.py") - thumbs = _reaction_steps(steps, "+1") - assert len(thumbs) == 1, "expected exactly one 👍 (+1) reaction step" - idx, step = thumbs[0] - assert idx > run_idx, "👍 must come AFTER the triage run" - assert "success()" in step["if"], "👍 must only fire when the reconsider run succeeded" - assert "vars.AGENT_SHIN_ENABLED == 'true'" in step["if"], ( - "👍 must be gated on AGENT_SHIN_ENABLED so dry-run stays inert" - ) From a4624b6c6c4bb1ddf753bd57bbfd37b711598c90 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:03:45 -0700 Subject: [PATCH 101/102] fix(gemini): derive the finish reason key set from the Candidates type Candidates.finishReason listed eleven values while the mapping key set carried twenty-one, so typed fixtures could not spell the reasons this PR handles. GeminiFinishReason is now the one list, the key set derives from it, and a test checks every documented reason has an explicit mapping instead of falling through to "stop" --- .../vertex_and_google_ai_studio_gemini.py | 29 ++------------ litellm/types/llms/vertex_ai.py | 39 ++++++++++++------- ...test_vertex_and_google_ai_studio_gemini.py | 9 ++++- 3 files changed, 36 insertions(+), 41 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a95b845718a..46f1b948026 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -6,7 +6,7 @@ import time from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args import httpx @@ -57,6 +57,7 @@ from litellm.types.llms.vertex_ai import ( ContentType, FunctionCallingConfig, FunctionDeclaration, + GeminiFinishReason, GeminiThinkingConfig, GenerateContentResponseBody, HttpxPartType, @@ -1330,31 +1331,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.", } - _GEMINI_FINISH_REASON_KEYS = frozenset( - { - "STOP", - "MAX_TOKENS", - "SAFETY", - "RECITATION", - "FINISH_REASON_UNSPECIFIED", - "MALFORMED_FUNCTION_CALL", - "LANGUAGE", - "OTHER", - "BLOCKLIST", - "PROHIBITED_CONTENT", - "SPII", - "IMAGE_SAFETY", - "IMAGE_PROHIBITED_CONTENT", - "TOO_MANY_TOOL_CALLS", - "MALFORMED_RESPONSE", - "NO_IMAGE", - "IMAGE_RECITATION", - "IMAGE_OTHER", - "ESCALATION", - "UNEXPECTED_TOOL_CALL", - "MISSING_THOUGHT_SIGNATURE", - } - ) + _GEMINI_FINISH_REASON_KEYS: Final[frozenset[str]] = frozenset(get_args(GeminiFinishReason)) @staticmethod def get_finish_reason_mapping() -> dict[str, OpenAIChatCompletionFinishReason]: diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 3b95b786631..ce51e46ef15 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -425,22 +425,35 @@ class UrlContextMetadata(TypedDict, total=False): urlMetadata: list[UrlMetadata] +GeminiFinishReason = Literal[ + "FINISH_REASON_UNSPECIFIED", + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "LANGUAGE", + "OTHER", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "SPII", + "MALFORMED_FUNCTION_CALL", + "IMAGE_SAFETY", + "IMAGE_PROHIBITED_CONTENT", + "TOO_MANY_TOOL_CALLS", + "MALFORMED_RESPONSE", + "NO_IMAGE", + "IMAGE_RECITATION", + "IMAGE_OTHER", + "ESCALATION", + "UNEXPECTED_TOOL_CALL", + "MISSING_THOUGHT_SIGNATURE", +] + + class Candidates(TypedDict, total=False): index: int content: HttpxContentType - finishReason: Literal[ - "FINISH_REASON_UNSPECIFIED", - "STOP", - "MAX_TOKENS", - "SAFETY", - "RECITATION", - "OTHER", - "BLOCKLIST", - "PROHIBITED_CONTENT", - "SPII", - "MALFORMED_FUNCTION_CALL", - "IMAGE_SAFETY", - ] + finishReason: GeminiFinishReason safetyRatings: list[SafetyRatings] citationMetadata: CitationMetadata groundingMetadata: GroundingMetadata diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 4ee199bd9af..6c818016c87 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2,7 +2,7 @@ import asyncio import json import re from copy import deepcopy -from typing import Final, List, cast +from typing import Final, List, cast, get_args from unittest.mock import MagicMock, patch import httpx @@ -18,7 +18,7 @@ from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) -from litellm.types.llms.vertex_ai import UsageMetadata +from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage from litellm.utils import CustomStreamWrapper @@ -940,6 +940,11 @@ def test_check_finish_reason(): ) +def test_every_documented_gemini_finish_reason_has_an_explicit_mapping(): + documented: Final = frozenset(get_args(GeminiFinishReason)) + assert set(VertexGeminiConfig.get_finish_reason_mapping()) == documented + + def test_finish_reason_unspecified_and_malformed_function_call(): """ Test that FINISH_REASON_UNSPECIFIED and MALFORMED_FUNCTION_CALL From 1704aeebb49eea499a863a2dfbd9c990b0e2b501 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 23:06:47 +0000 Subject: [PATCH 102/102] fix(enterprise): resolve openai_moderations model at call time and default to omni-moderation-latest Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../enterprise_hooks/openai_moderation.py | 9 +++-- litellm/constants.py | 2 ++ .../guardrails/test_guardrail_coverage.py | 36 +++++++++++++++++++ 3 files changed, 42 insertions(+), 5 deletions(-) diff --git a/enterprise/enterprise_hooks/openai_moderation.py b/enterprise/enterprise_hooks/openai_moderation.py index 2162370804a..017f51bfabd 100644 --- a/enterprise/enterprise_hooks/openai_moderation.py +++ b/enterprise/enterprise_hooks/openai_moderation.py @@ -17,6 +17,7 @@ from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_OPENAI_MODERATIONS_MODEL from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import iter_message_text @@ -24,11 +25,9 @@ from litellm.types.utils import CallTypesLiteral class _ENTERPRISE_OpenAI_Moderation(CustomLogger): - def __init__(self): - self.model_name = ( - litellm.openai_moderations_model_name or "text-moderation-latest" - ) # pass the model_name you initialized on litellm.Router() - pass + @property + def model_name(self) -> str: + return litellm.openai_moderations_model_name or DEFAULT_OPENAI_MODERATIONS_MODEL #### CALL HOOKS - proxy only #### diff --git a/litellm/constants.py b/litellm/constants.py index e7cb21a3a7d..a7d4eba0f15 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -158,6 +158,8 @@ DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL: Final = str( ) DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD", 0.75)) +DEFAULT_OPENAI_MODERATIONS_MODEL: Final = "omni-moderation-latest" + # MCP OAuth2 Client Credentials Defaults MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200")) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index f25e83b1672..548677c70bc 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -18,7 +18,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from httpx import Request, Response +import litellm from litellm import DualCache +from litellm.constants import DEFAULT_OPENAI_MODERATIONS_MODEL from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import Choices, Message, ModelResponse @@ -764,6 +766,40 @@ async def test_openai_moderation_inspects_multimodal_content(monkeypatch, user_a assert seen_inputs == ["alpha beta"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("configured_after_init", "expected_model"), + [("omni-moderation-2024-09-26", "omni-moderation-2024-09-26"), (None, DEFAULT_OPENAI_MODERATIONS_MODEL)], +) +async def test_openai_moderation_reads_model_name_at_call_time( + monkeypatch, user_api_key, configured_after_init, expected_model +): + """``litellm_settings`` applies ``callbacks`` and ``openai_moderations_model_name`` in YAML + order, so the hook must resolve the model when it runs, not when it is constructed.""" + from enterprise.enterprise_hooks.openai_moderation import ( + _ENTERPRISE_OpenAI_Moderation, + ) + + monkeypatch.setattr(litellm, "openai_moderations_model_name", None) + guard = _ENTERPRISE_OpenAI_Moderation() + monkeypatch.setattr(litellm, "openai_moderations_model_name", configured_after_init) + + class FakeModeration: + results = [type("R", (), {"flagged": False})()] + + fake_router = MagicMock() + fake_router.amoderation = AsyncMock(return_value=FakeModeration()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router, raising=False) + + await guard.async_moderation_hook( + data={"messages": [{"role": "user", "content": "hello"}]}, + user_api_key_dict=user_api_key, + call_type="acompletion", + ) + + fake_router.amoderation.assert_awaited_once_with(model=expected_model, input="hello") + + # ── Google Text Moderation ────────────────────────────────────────────────────