diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
index 973b2ecee80..e4bf67ab72f 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
@@ -2047,6 +2047,82 @@ describe("TeamInfoView", () => {
expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled();
});
+ it("saves a team admin's metadata edit with the proxy's own keys carried over", async () => {
+ const user = userEvent.setup({ delay: null });
+ vi.mocked(useTeamMetadataSchema).mockReturnValue({
+ data: [{ key: "cost_center", label: "Cost Center" }],
+ isLoading: false,
+ } as any);
+ vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
+ vi.mocked(networking.teamInfoCall).mockResolvedValue(
+ createMockTeamData({
+ metadata: { cost_center: "eng", priority: "high" },
+ caller_edit_access: {
+ kind: "team_admin",
+ editable_fields: ["metadata"],
+ proxy_controlled_metadata_keys: ["priority"],
+ },
+ }),
+ );
+
+ renderWithProviders(
);
+
+ await user.click(await screen.findByRole("tab", { name: "Settings" }));
+ await user.click(await screen.findByRole("button", { name: /edit settings/i }));
+ const costCenter = await screen.findByLabelText("Cost Center");
+ expect(costCenter).toHaveValue("eng");
+ expect(screen.queryByDisplayValue("priority")).not.toBeInTheDocument();
+
+ await user.clear(costCenter);
+ await user.type(costCenter, "sales");
+ await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+ await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
+ expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1]).toMatchObject({
+ metadata: { cost_center: "sales", priority: "high" },
+ });
+ });
+
+ it("saves other fields when a declared metadata field is also a proxy setting", async () => {
+ const user = userEvent.setup({ delay: null });
+ vi.mocked(useTeamMetadataSchema).mockReturnValue({
+ data: [
+ { key: "cost_center", label: "Cost Center" },
+ { key: "priority", label: "Priority" },
+ ],
+ isLoading: false,
+ } as any);
+ vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
+ vi.mocked(networking.teamInfoCall).mockResolvedValue(
+ createMockTeamData({
+ tpm_limit: 1000,
+ metadata: { cost_center: "eng", priority: "high" },
+ caller_edit_access: {
+ kind: "team_admin",
+ editable_fields: ["metadata", "tpm_limit"],
+ proxy_controlled_metadata_keys: ["priority"],
+ },
+ }),
+ );
+
+ renderWithProviders(
);
+
+ await user.click(await screen.findByRole("tab", { name: "Settings" }));
+ await user.click(await screen.findByRole("button", { name: /edit settings/i }));
+ expect(screen.queryByLabelText("Priority")).not.toBeInTheDocument();
+ expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled();
+
+ const tpm = screen.getByLabelText("Tokens per minute Limit (TPM)");
+ await user.clear(tpm);
+ await user.type(tpm, "2000");
+ await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+ await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
+ const sent = vi.mocked(networking.teamUpdateCall).mock.calls[0][1];
+ expect(sent).toMatchObject({ tpm_limit: 2000 });
+ expect(sent).not.toHaveProperty("metadata");
+ });
+
it("prefills the member defaults a team admin may edit from the team's member budget", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
index 24d102aebf7..29c048034f4 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -56,6 +56,8 @@ import {
parseTeamEditAccess,
TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION,
TEAM_ADMIN_EDITING_DISABLED_TITLE,
+ teamAdminDeclaredFields,
+ teamAdminMetadataPairs,
type TeamAdminSettingsChanges,
} from "./teamAdminEditAccess";
import TeamAdminSettingsForm from "./TeamAdminSettingsForm";
@@ -1305,6 +1307,11 @@ const TeamInfoView: React.FC
= ({
const { team_info: info } = teamData;
+ const teamAdminSchemaFields =
+ teamEditAccess.kind === "team_admin"
+ ? teamAdminDeclaredFields(teamMetadataSchemaFields, teamEditAccess.proxyControlledMetadataKeys)
+ : [];
+ const teamAdminDeclaredKeys = teamAdminSchemaFields.map((field) => field.key);
const teamAdminSettingsEditor =
teamEditAccess.kind === "team_admin" ? (
= ({
team_member_tpm_limit: info.team_member_budget_table?.tpm_limit,
team_member_rpm_limit: info.team_member_budget_table?.rpm_limit,
}}
+ metadata={{
+ pairs: teamAdminMetadataPairs(
+ info.metadata,
+ teamEditAccess.proxyControlledMetadataKeys,
+ teamAdminDeclaredKeys,
+ ),
+ context: {
+ storedMetadata: info.metadata ?? {},
+ proxyControlledKeys: teamEditAccess.proxyControlledMetadataKeys,
+ declaredKeys: new Set(teamAdminDeclaredKeys),
+ },
+ schemaFields: teamAdminSchemaFields,
+ schemaLoading: isTeamMetadataSchemaLoading,
+ }}
editableFields={teamEditAccess.editableFields}
mayRaiseMaxBudget={teamEditAccess.mayRaiseMaxBudget}
isSaving={isTeamSaving}
diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts
index 84898563899..731755959e1 100644
--- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts
+++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts
@@ -1,6 +1,11 @@
import { describe, expect, it } from "vitest";
import {
+ type TeamAdminMetadataContext,
+ teamAdminDeclaredFields,
+ teamAdminMetadataChange,
+ teamAdminMetadataPairIssue,
+ teamAdminMetadataPairs,
parseSupportedTeamAdminEditableFields,
parseTeamAdminEditableFields,
parseTeamEditAccess,
@@ -14,6 +19,7 @@ describe("teamAdminFieldLabel", () => {
["rpm_limit", "Requests per minute Limit (RPM)"],
["max_budget", "Max Budget (USD)"],
["soft_budget", "Soft Budget (USD)"],
+ ["metadata", "Team Metadata"],
["raise_max_budget", "Raise the team's max budget"],
["projects", "Create and update projects"],
["team_member_budget", "Member Default Budget (USD)"],
@@ -213,11 +219,14 @@ describe("parseTeamEditAccess", () => {
});
it("hands a team admin the fields the proxy enabled", () => {
- expect(parseTeamEditAccess({ kind: "team_admin", editable_fields: ["tpm_limit"] })).toEqual({
+ const expected = {
kind: "team_admin",
editableFields: new Set(["tpm_limit"]),
mayRaiseMaxBudget: false,
- });
+ proxyControlledMetadataKeys: new Set(),
+ };
+
+ expect(parseTeamEditAccess({ kind: "team_admin", editable_fields: ["tpm_limit"] })).toEqual(expected);
});
it.each([
@@ -226,10 +235,26 @@ describe("parseTeamEditAccess", () => {
["an older proxy omits the flag", {}, false],
["the flag is not a boolean", { may_raise_max_budget: "yes" }, false],
])("lets a team admin raise the max budget only when %s", (_label, flag, mayRaiseMaxBudget) => {
- expect(parseTeamEditAccess({ kind: "team_admin", editable_fields: ["max_budget"], ...flag })).toEqual({
+ const expected = {
kind: "team_admin",
editableFields: new Set(["max_budget"]),
mayRaiseMaxBudget,
+ proxyControlledMetadataKeys: new Set(),
+ };
+
+ expect(parseTeamEditAccess({ kind: "team_admin", editable_fields: ["max_budget"], ...flag })).toEqual(expected);
+ });
+
+ it("hands a team admin the metadata keys the proxy keeps for proxy admins", () => {
+ const access = parseTeamEditAccess({
+ kind: "team_admin",
+ editable_fields: ["metadata"],
+ proxy_controlled_metadata_keys: ["guardrails", "priority"],
+ });
+
+ expect(access).toMatchObject({
+ kind: "team_admin",
+ proxyControlledMetadataKeys: new Set(["guardrails", "priority"]),
});
});
@@ -242,3 +267,101 @@ describe("parseTeamEditAccess", () => {
expect(parseTeamEditAccess(value)).toEqual({ kind: "none" });
});
});
+
+describe("teamAdminMetadataChange", () => {
+ const context: TeamAdminMetadataContext = {
+ storedMetadata: { cost_center: "eng", region: "eu", priority: "high", team_member_budget_id: "budget-1" },
+ proxyControlledKeys: new Set(["priority", "team_member_budget_id", "guardrails"]),
+ declaredKeys: new Set(["cost_center", "department"]),
+ };
+ const storedPairs = [
+ { key: "cost_center", value: "eng" },
+ { key: "region", value: "eu" },
+ ];
+
+ it("sends nothing when the editable keys are unchanged", () => {
+ expect(teamAdminMetadataChange(storedPairs, context)).toStrictEqual({});
+ });
+
+ it("sends the full metadata with the proxy settings carried over when a key is added", () => {
+ expect(teamAdminMetadataChange([...storedPairs, { key: "project", value: "atlas" }], context)).toStrictEqual({
+ metadata: {
+ cost_center: "eng",
+ region: "eu",
+ project: "atlas",
+ priority: "high",
+ team_member_budget_id: "budget-1",
+ },
+ });
+ });
+
+ it("sends the full metadata without a deleted key", () => {
+ expect(teamAdminMetadataChange([{ key: "cost_center", value: "eng" }], context)).toStrictEqual({
+ metadata: { cost_center: "eng", priority: "high", team_member_budget_id: "budget-1" },
+ });
+ });
+
+ it("leaves out an empty declared field the team never set", () => {
+ expect(teamAdminMetadataChange([...storedPairs, { key: "department", value: "" }], context)).toStrictEqual({});
+ });
+});
+
+describe("teamAdminMetadataPairIssue", () => {
+ const context: TeamAdminMetadataContext = {
+ storedMetadata: { cost_center: "eng" },
+ proxyControlledKeys: new Set(["priority"]),
+ declaredKeys: new Set(["cost_center", "department"]),
+ };
+
+ it("flags a key the proxy keeps for proxy admins", () => {
+ expect(teamAdminMetadataPairIssue({ key: "priority", value: "low" }, context)?.path).toBe("key");
+ });
+
+ it("flags emptying a declared field the team has set", () => {
+ expect(teamAdminMetadataPairIssue({ key: "cost_center", value: "" }, context)?.path).toBe("value");
+ });
+
+ it.each([
+ ["JSON null", "null"],
+ ["a JSON empty string", '""'],
+ ["only spaces", " "],
+ ])("flags a declared field the team has set changed to %s", (_label, value) => {
+ expect(teamAdminMetadataPairIssue({ key: "cost_center", value }, context)?.path).toBe("value");
+ });
+
+ it("accepts a declared field the proxy admin stored empty, resent as it is", () => {
+ const storedEmpty = { ...context, storedMetadata: { cost_center: "" } };
+
+ expect(teamAdminMetadataPairIssue({ key: "cost_center", value: "" }, storedEmpty)).toBeNull();
+ });
+
+ it.each([
+ ["an ordinary key", { key: "project", value: "" }],
+ ["a declared field the team never set", { key: "department", value: "" }],
+ ["a declared field with a new value", { key: "cost_center", value: "sales" }],
+ ])("accepts %s", (_label, pair) => {
+ expect(teamAdminMetadataPairIssue(pair, context)).toBeNull();
+ });
+});
+
+describe("teamAdminMetadataPairs", () => {
+ it("lists declared fields first in schema order, then the team's own keys, without the proxy's keys", () => {
+ const metadata = { region: "eu", priority: "high", department: "r&d", cost_center: "eng" };
+ const pairs = teamAdminMetadataPairs(metadata, new Set(["priority"]), ["cost_center", "department"]);
+
+ expect(pairs.map((pair) => pair.key)).toStrictEqual(["cost_center", "department", "region"]);
+ });
+});
+
+describe("teamAdminDeclaredFields", () => {
+ it("leaves out declared fields that are also proxy settings", () => {
+ const fields = [
+ { key: "cost_center", label: "Cost Center" },
+ { key: "priority", label: "Priority" },
+ ];
+
+ expect(teamAdminDeclaredFields(fields, new Set(["priority"]))).toStrictEqual([
+ { key: "cost_center", label: "Cost Center" },
+ ]);
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
index 28ba5eed3e4..bf8d56b6110 100644
--- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
+++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
@@ -1,23 +1,34 @@
import { z } from "zod";
import { numberOrNull } from "@/lib/forms/numberOrNull";
+import {
+ type MetadataPair,
+ metadataObjectToPairs,
+ metadataPairsToObject,
+} from "../common_components/MetadataKeyValueFields";
+
export const TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING = "team_admin_editable_team_fields";
export const TEAM_ADMIN_RAISE_MAX_BUDGET_PERMISSION = "raise_max_budget";
export const TEAM_ADMIN_CREATE_MEMBER_KEYS_PERMISSION = "create_member_keys";
+export const TEAM_ADMIN_METADATA_FIELD = "metadata";
+
export const TEAM_ADMIN_EDITING_DISABLED_TITLE = "Team admins cannot edit team settings on this proxy";
export const TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION =
"Ask a proxy admin to enable fields under Settings > UI > Team admin editable fields.";
+const teamAdminEditAccessShape = {
+ kind: z.literal("team_admin"),
+ editable_fields: z.array(z.string()),
+ may_raise_max_budget: z.boolean().catch(false),
+ proxy_controlled_metadata_keys: z.array(z.string()).catch([]),
+};
+
const callerEditAccessSchema = z.discriminatedUnion("kind", [
z.object({ kind: z.literal("unrestricted") }),
- z.object({
- kind: z.literal("team_admin"),
- editable_fields: z.array(z.string()),
- may_raise_max_budget: z.boolean().catch(false),
- }),
+ z.object(teamAdminEditAccessShape),
z.object({ kind: z.literal("team_admin_disabled") }),
z.object({ kind: z.literal("none") }),
]);
@@ -26,7 +37,12 @@ export type CallerEditAccess = z.infer;
export type TeamEditAccess =
| { readonly kind: "unrestricted" }
- | { readonly kind: "team_admin"; readonly editableFields: ReadonlySet; readonly mayRaiseMaxBudget: boolean }
+ | {
+ readonly kind: "team_admin";
+ readonly editableFields: ReadonlySet;
+ readonly mayRaiseMaxBudget: boolean;
+ readonly proxyControlledMetadataKeys: ReadonlySet;
+ }
| { readonly kind: "team_admin_disabled" }
| { readonly kind: "none" };
@@ -73,6 +89,7 @@ const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([
["tpm_limit", "Tokens per minute Limit (TPM)"],
["rpm_limit", "Requests per minute Limit (RPM)"],
["max_budget", "Max Budget (USD)"],
+ [TEAM_ADMIN_METADATA_FIELD, "Team Metadata"],
["soft_budget", "Soft Budget (USD)"],
["team_member_budget", "Member Default Budget (USD)"],
["team_member_budget_duration", "Member Default Budget Duration"],
@@ -86,11 +103,13 @@ const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([
export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_LABELS.get(field) ?? field;
-export type TeamAdminSettingsValues = { readonly [F in TeamAdminSettingsField]?: string | number | null };
+export type TeamAdminSettingsValues = { readonly [F in TeamAdminNumericSettingsField]?: string | number | null } & {
+ readonly [F in TeamAdminDurationSettingsField]?: string | null;
+};
export type TeamAdminSettingsChanges = { readonly [F in TeamAdminNumericSettingsField]?: number | null } & {
readonly [F in TeamAdminDurationSettingsField]?: string | null;
-};
+} & { readonly metadata?: Record };
const durationOrNull = (value: string | number | null | undefined): string | null =>
typeof value === "string" && value !== "" ? value : null;
@@ -129,7 +148,81 @@ export const parseTeamEditAccess = (callerEditAccess: unknown): TeamEditAccess =
kind: "team_admin",
editableFields: new Set(parsed.data.editable_fields),
mayRaiseMaxBudget: parsed.data.may_raise_max_budget,
+ proxyControlledMetadataKeys: new Set(parsed.data.proxy_controlled_metadata_keys),
};
}
return parsed.data;
};
+
+export interface TeamAdminMetadataContext {
+ readonly storedMetadata: Readonly>;
+ readonly proxyControlledKeys: ReadonlySet;
+ readonly declaredKeys: ReadonlySet;
+}
+
+const canonicalJson = (value: unknown): string =>
+ JSON.stringify(value, (_key, nested: unknown) =>
+ nested !== null && typeof nested === "object" && !Array.isArray(nested)
+ ? Object.fromEntries(Object.entries(nested).sort(([a], [b]) => a.localeCompare(b)))
+ : nested,
+ );
+
+const pairValue = (pair: Partial | undefined): unknown =>
+ metadataPairsToObject([{ key: "value", value: pair?.value ?? "" }]).value;
+
+const isBlankString = (value: unknown): boolean => typeof value === "string" && value.trim() === "";
+
+const isBlank = (value: unknown): boolean => value === null || value === undefined || isBlankString(value);
+
+const isUnsetDeclaredPair = (pair: Partial | undefined, context: TeamAdminMetadataContext): boolean => {
+ const key = pair?.key;
+ if (key === undefined || !context.declaredKeys.has(key)) return false;
+ return !(key in context.storedMetadata) && isBlank(pairValue(pair));
+};
+
+export const teamAdminDeclaredFields = (
+ schemaFields: readonly Field[],
+ proxyControlledKeys: ReadonlySet,
+): Field[] => schemaFields.filter((field) => !proxyControlledKeys.has(field.key));
+
+export const teamAdminMetadataPairs = (
+ metadata: Readonly> | null | undefined,
+ proxyControlledKeys: ReadonlySet,
+ declaredKeys: readonly string[],
+): MetadataPair[] => {
+ const pairs = metadataObjectToPairs(metadata, proxyControlledKeys);
+ const declaredRank = (key: string): number => {
+ const rank = declaredKeys.indexOf(key);
+ return rank === -1 ? declaredKeys.length : rank;
+ };
+ return pairs.toSorted((a, b) => declaredRank(a.key) - declaredRank(b.key));
+};
+
+export const teamAdminMetadataChange = (
+ pairs: readonly (Partial | undefined)[] | undefined,
+ context: TeamAdminMetadataContext,
+): { readonly metadata?: Record } => {
+ const edited = metadataPairsToObject((pairs ?? []).filter((pair) => !isUnsetDeclaredPair(pair, context)));
+ const proxyControlled = Object.entries(context.storedMetadata).filter(([key]) =>
+ context.proxyControlledKeys.has(key),
+ );
+ const metadata = { ...Object.fromEntries(proxyControlled), ...edited };
+ return canonicalJson(metadata) === canonicalJson(context.storedMetadata) ? {} : { metadata };
+};
+
+export interface TeamAdminMetadataPairIssue {
+ readonly path: "key" | "value";
+ readonly message: string;
+}
+
+export const teamAdminMetadataPairIssue = (
+ pair: Partial | undefined,
+ context: TeamAdminMetadataContext,
+): TeamAdminMetadataPairIssue | null => {
+ if (pair?.key === undefined) return null;
+ if (context.proxyControlledKeys.has(pair.key)) return { path: "key", message: "Only proxy admins can set this key" };
+ if (!context.declaredKeys.has(pair.key) || !(pair.key in context.storedMetadata)) return null;
+ const value = pairValue(pair);
+ const clearsStoredValue = isBlank(value) && canonicalJson(value) !== canonicalJson(context.storedMetadata[pair.key]);
+ return clearsStoredValue ? { path: "value", message: "This field cannot be left empty" } : null;
+};
From 6d1a3bce2c3996d226434d8cf64cf26927044ba8 Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 16:58:00 -0700
Subject: [PATCH 058/140] fix(responses): return the first message item as
content for structured output on the chat bridge (#46508)
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../handler.py | 20 +++-
.../transformation.py | 89 +++++++++++++--
.../test_responses_bridge_incomplete.py | 105 ++++++++++++++++++
...responses_transformation_transformation.py | 91 +++++++++++++++
4 files changed, 293 insertions(+), 12 deletions(-)
diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py
index edc32dc487a..3f7017eccc2 100644
--- a/litellm/completion_extras/litellm_responses_transformation/handler.py
+++ b/litellm/completion_extras/litellm_responses_transformation/handler.py
@@ -2,8 +2,14 @@
Handler for transforming /chat/completions api requests to litellm.responses requests
"""
-from collections.abc import AsyncIterable, Coroutine, Iterable
-from typing import TYPE_CHECKING, Any, Final, Union
+from collections.abc import AsyncIterable, Coroutine, Iterable, Mapping
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Final,
+ Union,
+ cast, # noqa: TID251 # optional_params reaches the bridge as an untyped dict
+)
from typing_extensions import TypedDict
@@ -245,10 +251,15 @@ class ResponsesToCompletionBridgeHandler:
else:
if self._is_preformatted_cached_chat_stream(result):
return self._apply_post_stream_processing(result, model, custom_llm_provider)
+ from .transformation import is_structured_output_request
+
completion_stream: Final = self.transformation_handler.get_model_response_iterator(
streaming_response=result,
sync_stream=True,
json_mode=kwargs.get("json_mode"),
+ structured_output=is_structured_output_request(
+ cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
+ ),
)
streamwrapper: Final = CustomStreamWrapper(
completion_stream=completion_stream,
@@ -343,10 +354,15 @@ class ResponsesToCompletionBridgeHandler:
else:
if self._is_preformatted_cached_chat_stream(result):
return self._apply_post_stream_processing(result, model, custom_llm_provider)
+ from .transformation import is_structured_output_request
+
completion_stream: Final = self.transformation_handler.get_model_response_iterator(
streaming_response=result,
sync_stream=False,
json_mode=kwargs.get("json_mode"),
+ structured_output=is_structured_output_request(
+ cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
+ ),
)
streamwrapper: Final = CustomStreamWrapper(
completion_stream=completion_stream,
diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py
index 2e4bb47654e..638acfe0723 100644
--- a/litellm/completion_extras/litellm_responses_transformation/transformation.py
+++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py
@@ -20,7 +20,7 @@ from openai.types.responses.response_input_param import (
from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam
from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam
from openai.types.responses.tool_param import FunctionToolParam
-from pydantic import BaseModel, TypeAdapter, ValidationError
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm import ModelResponse
@@ -346,14 +346,20 @@ def _fold_choices_into_one(
tool_calls: Sequence[Mapping[str, object]],
trailing_reasoning_content: str | None,
trailing_reasoning_item: _BuiltReasoningItem | None,
+ structured_output: bool = False,
) -> "Choices":
- """Clients that only read choices[0] must still see every text and tool call."""
+ """Clients that only read choices[0] must still see every text and tool call.
+
+ Under a structured ``response_format`` each message item is a complete document, so
+ ``content`` is the first one, the same item the openai SDK returns as ``output_parsed``.
+ """
from litellm.types.utils import Choices, Message
messages: Final = tuple(choice.message for choice in choices)
- texts: Final = tuple(str(message.content or "") for message in messages)
+ content_messages: Final = messages[:1] if structured_output else messages
+ texts: Final = tuple(str(message.content or "") for message in content_messages)
offsets: Final = (0, *accumulate(len(text) for text in texts))
- annotations: Final = tuple(_shifted_annotations(messages, offsets))
+ annotations: Final = tuple(_shifted_annotations(content_messages, offsets))
reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content)
reasoning_items: Final = (
*chain.from_iterable(_message_reasoning_items(m) for m in messages),
@@ -373,6 +379,38 @@ def _fold_choices_into_one(
)
+def _responses_event_type(chunk: Mapping[str, object]) -> object:
+ event_type: Final = chunk.get("type")
+ return event_type.value if isinstance(event_type, ResponsesAPIStreamEvents) else event_type
+
+
+class _TypedPayload(BaseModel):
+ """The ``type`` discriminator of a response_format dict or a Responses output item."""
+
+ model_config = ConfigDict(frozen=True, extra="ignore")
+ type: str | None = None
+
+
+_TYPED_PAYLOAD: Final = TypeAdapter(_TypedPayload)
+
+
+def _payload_type(payload: object) -> str | None:
+ try:
+ return _TYPED_PAYLOAD.validate_python(payload).type
+ except ValidationError:
+ return None
+
+
+def _is_message_item_done(chunk: Mapping[str, object]) -> bool:
+ if _responses_event_type(chunk) != ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
+ return False
+ return _payload_type(chunk.get("item")) == "message"
+
+
+def is_structured_output_request(optional_params: Mapping[str, object]) -> bool:
+ return _payload_type(optional_params.get("response_format")) in ("json_schema", "json_object")
+
+
_ToolChoiceT = TypeVar("_ToolChoiceT")
@@ -897,6 +935,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def _convert_response_output_to_choices(
output_items: Sequence[object],
handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None,
+ structured_output: bool = False,
) -> list["Choices"]:
"""
Convert Responses API output items to chat completion choices.
@@ -904,6 +943,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
Args:
output_items: List of items from ResponsesAPIResponse.output
handle_raw_dict_callback: Optional callback for handling raw dict items
+ structured_output: The chat request carried a json response_format, so only the
+ first message item becomes content
Returns:
List of Choices objects
@@ -1027,6 +1068,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
tool_calls=accumulated_tool_calls,
trailing_reasoning_content=reasoning_content,
trailing_reasoning_item=pending_reasoning_item,
+ structured_output=structured_output,
)
]
@@ -1157,6 +1199,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
choices: Final = self._convert_response_output_to_choices(
output_items=output_items,
handle_raw_dict_callback=self._handle_raw_dict_response_item,
+ structured_output=is_structured_output_request(
+ cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict
+ ),
)
response_is_incomplete: Final = raw_response.status == "incomplete" or (
@@ -1219,8 +1264,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
sync_stream: bool,
json_mode: bool | None = False,
+ structured_output: bool = False,
) -> BaseModelResponseIterator:
- return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode)
+ return OpenAiResponsesToChatCompletionStreamIterator(
+ streaming_response, sync_stream, json_mode, structured_output=structured_output
+ )
def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]:
if role in ("user", "system", "developer", "tool"):
@@ -1587,11 +1635,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
sync_stream: bool,
json_mode: bool | None = False,
+ structured_output: bool = False,
):
super().__init__(streaming_response, sync_stream, json_mode)
self._chat_completion_id: str | None = None
self._served_service_tier: str | None = None
self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state
+ self._structured_output: Final = structured_output
+ self._first_message_item_done: bool = False
def _handle_string_chunk(
self, str_line: Union[str, "BaseModel"]
@@ -1903,13 +1954,31 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
"""
verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk)
self._remember_served_service_tier(chunk)
- return self._with_served_service_tier(
- self._with_stream_scoped_id(
- OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
- chunk, tool_call_index_map=self._tool_call_index_map
- )
+ translated: Final = (
+ self._empty_content_chunk()
+ if self._is_text_delta_after_structured_answer(chunk)
+ else OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
+ chunk, tool_call_index_map=self._tool_call_index_map
)
)
+ if self._structured_output and _is_message_item_done(chunk):
+ self._first_message_item_done = True
+ return self._with_served_service_tier(self._with_stream_scoped_id(translated))
+
+ def _is_text_delta_after_structured_answer(self, chunk: Mapping[str, object]) -> bool:
+ """Under a structured response_format the first message item is the whole answer, so
+ text from any later message item is dropped, matching the non-stream fold."""
+ return (
+ self._structured_output
+ and self._first_message_item_done
+ and _responses_event_type(chunk) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA.value
+ )
+
+ @staticmethod
+ def _empty_content_chunk() -> "ModelResponseStream":
+ from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+ return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason=None)])
def _remember_served_service_tier(self, chunk: dict[str, object]) -> None:
response_payload: Final = chunk.get("response")
diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py
index e970c806330..e3978b4debd 100644
--- a/tests/integration/providers/test_responses_bridge_incomplete.py
+++ b/tests/integration/providers/test_responses_bridge_incomplete.py
@@ -1297,3 +1297,108 @@ def test_chat_stream_over_responses_deployment_folds_message_and_two_function_ca
assert arguments_by_index == {0: '{"city":"Paris"}', 1: '{"city":"Paris","tz":"CET"}'}, choices
assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["tool_calls"], choices
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")]
+
+
+_TWO_FINAL_ANSWERS: Final = (
+ _message_item("msg_final_first", '{"city":"Paris"}', phase="final_answer"),
+ _message_item("msg_final_second", '{"city":"Paris"}', phase="final_answer"),
+)
+
+_CITY_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "answer",
+ "strict": True,
+ "schema": {
+ "type": "object",
+ "properties": {"city": {"type": "string"}},
+ "required": ["city"],
+ "additionalProperties": False,
+ },
+ },
+}
+
+
+def test_chat_over_responses_deployment_with_json_schema_returns_first_of_two_final_answers(
+ gateway: Gateway,
+) -> None:
+ """gpt-5.4-mini answered a strict json_schema request on /responses with two final_answer items,
+ each the whole document (captured 2026-10-14); a chat client must get one parseable document."""
+ identity: Final = "responses-bridge-json-schema-two-finals-" + uuid.uuid4().hex
+
+ def respond(request: Request) -> Reply:
+ if request.method == "GET" and request.target == "/v1/models":
+ return Reply(body=b'{"object":"list","data":[]}')
+ assert request.method == "POST" and request.target == "/responses", request.target
+ assert json.loads(request.body)["text"]["format"]["type"] == "json_schema", request.body
+ return Reply(body=_responses_body("resp_two_finals", _TWO_FINAL_ANSWERS))
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model="openai/responses/gpt-5.4-mini", api_base=wire.url, api_key="synthetic-openai-key"
+ )
+ response: Final = gateway.request(
+ "POST",
+ "/v1/chat/completions",
+ {
+ "model": model,
+ "messages": [{"role": "user", "content": f"What is the capital of France? {identity}"}],
+ "response_format": _CITY_RESPONSE_FORMAT,
+ "cache": {"no-cache": True},
+ },
+ )
+ assert response.status_code == 200, response.text
+ body: Final = response.json()
+ assert body["choices"] == [
+ {
+ "finish_reason": "stop",
+ "index": 0,
+ "message": {"role": "assistant", "content": '{"city":"Paris"}'},
+ }
+ ], response.text
+ assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")]
+
+
+def test_chat_stream_over_responses_deployment_with_json_schema_streams_first_of_two_final_answers(
+ gateway: Gateway,
+) -> None:
+ identity: Final = "responses-bridge-json-schema-two-finals-stream-" + uuid.uuid4().hex
+
+ def respond(request: Request) -> Reply:
+ if request.method == "GET" and request.target == "/v1/models":
+ return Reply(body=b'{"object":"list","data":[]}')
+ assert request.method == "POST" and request.target == "/responses", request.target
+ upstream_body: Final = json.loads(request.body)
+ assert upstream_body["stream"] is True, request.body
+ assert upstream_body["text"]["format"]["type"] == "json_schema", request.body
+ return _stream_reply(
+ "resp_two_finals_stream",
+ _TWO_FINAL_ANSWERS,
+ (
+ *_message_item_frames(0, _TWO_FINAL_ANSWERS[0]),
+ *_message_item_frames(1, _TWO_FINAL_ANSWERS[1]),
+ ),
+ )
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model="openai/responses/gpt-5.4-mini", api_base=wire.url, api_key="synthetic-openai-key"
+ )
+ response: Final = gateway.request(
+ "POST",
+ "/v1/chat/completions",
+ {
+ "model": model,
+ "messages": [{"role": "user", "content": f"What is the capital of France? {identity}"}],
+ "response_format": _CITY_RESPONSE_FORMAT,
+ "stream": True,
+ "cache": {"no-cache": True},
+ },
+ )
+ assert response.status_code == 200, response.text
+ choices: Final = _stream_choices(response)
+ assert [choice["index"] for choice in choices] == [0] * len(choices), choices
+ deltas: Final = [object_value(choice["delta"]) for choice in choices]
+ assert "".join(str(delta["content"]) for delta in deltas if delta.get("content")) == '{"city":"Paris"}', choices
+ assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["stop"], choices
+ assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")]
diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
index 8072919ed87..91bfb8d78d6 100644
--- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
+++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
@@ -5267,3 +5267,94 @@ def test_convert_response_output_numbers_two_function_calls_zero_and_one() -> No
tool_calls: Final = choices[0].message.tool_calls
assert tool_calls is not None
assert [(call.index, call.function.name) for call in tool_calls] == [(0, "get_weather"), (1, "get_time")]
+
+
+def _final_answer_message(item_id: str, text: str) -> ResponseOutputMessage:
+ return ResponseOutputMessage.model_validate(
+ {
+ "id": item_id,
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "phase": "final_answer",
+ "content": [{"type": "output_text", "text": text, "annotations": []}],
+ }
+ )
+
+
+@pytest.mark.parametrize(
+ "structured_output, expected_content",
+ [
+ (True, '{"city":"Paris"}'),
+ (False, '{"city":"Paris"}{"city":"Paris"}'),
+ ],
+)
+def test_convert_response_output_structured_output_keeps_first_message_item(
+ structured_output: bool, expected_content: str
+) -> None:
+ """gpt-5.4-mini on /responses with a strict json_schema and a 'say what you will do first'
+ system prompt returned two final_answer items, each the whole document (captured 2026-10-14)."""
+ choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( # pyright: ignore[reportPrivateUsage] # exercises the converter directly
+ (_final_answer_message("msg_1", '{"city":"Paris"}'), _final_answer_message("msg_2", '{"city":"Paris"}')),
+ structured_output=structured_output,
+ )
+
+ assert len(choices) == 1
+ assert choices[0].index == 0
+ assert choices[0].finish_reason == "stop"
+ assert choices[0].message.content == expected_content
+ if structured_output:
+ assert json.loads(expected_content) == {"city": "Paris"}
+
+
+@pytest.mark.parametrize(
+ "optional_params, expected",
+ [
+ ({"response_format": {"type": "json_schema", "json_schema": {"name": "a", "schema": {}}}}, True),
+ ({"response_format": {"type": "json_object"}}, True),
+ ({"response_format": {"type": "text"}}, False),
+ ({}, False),
+ ],
+)
+def test_is_structured_output_request(optional_params: dict[str, object], expected: bool) -> None:
+ from litellm.completion_extras.litellm_responses_transformation.transformation import (
+ is_structured_output_request,
+ )
+
+ assert is_structured_output_request(optional_params) is expected
+
+
+@pytest.mark.parametrize(
+ "structured_output, expected_content",
+ [
+ (True, '{"city":"Paris"}'),
+ (False, '{"city":"Paris"}{"city":"Paris"}'),
+ ],
+)
+def test_stream_structured_output_drops_text_after_first_message_item(
+ structured_output: bool, expected_content: str
+) -> None:
+ from litellm.completion_extras.litellm_responses_transformation.transformation import (
+ OpenAiResponsesToChatCompletionStreamIterator,
+ )
+
+ iterator = OpenAiResponsesToChatCompletionStreamIterator(
+ streaming_response=iter(()), sync_stream=True, structured_output=structured_output
+ )
+ events: Final[tuple[dict[str, object], ...]] = (
+ {"type": "response.output_item.added", "output_index": 0, "item": {"type": "message", "content": []}},
+ {"type": "response.output_text.delta", "output_index": 0, "delta": '{"city":'},
+ {"type": "response.output_text.delta", "output_index": 0, "delta": '"Paris"}'},
+ {"type": "response.output_item.done", "output_index": 0, "item": {"type": "message", "content": []}},
+ {"type": "response.output_item.added", "output_index": 1, "item": {"type": "message", "content": []}},
+ {"type": "response.output_text.delta", "output_index": 1, "delta": '{"city":'},
+ {"type": "response.output_text.delta", "output_index": 1, "delta": '"Paris"}'},
+ {"type": "response.output_item.done", "output_index": 1, "item": {"type": "message", "content": []}},
+ {"type": "response.completed"},
+ )
+
+ chunks: Final = [iterator.chunk_parser(event) for event in events]
+
+ content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks)
+ assert content == expected_content
+ assert chunks[-1].choices[0].finish_reason == "stop"
From 44e94eb9858d30ff26cc991793632122cc078279 Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 17:02:00 -0700
Subject: [PATCH 059/140] fix(cost): split xai spend rows across input and
output instead of booking the total as output (#46320)
* fix(cost): keep the input/output split on streamed xai spend rows
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(cost): split xAI's reported total across input and output by token math
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../streaming_chunk_builder_utils.py | 14 +++++--
litellm/llms/xai/cost_calculator.py | 18 +++++++-
.../integration/cost_calculation/cases/xai.py | 10 ++++-
.../unit/llms/xai/test_xai_cost_calculator.py | 14 ++++---
tests/unit/test_main.py | 42 +++++++++++++++++++
5 files changed, 86 insertions(+), 12 deletions(-)
diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py
index 2ceb0042f1f..9dde1432fae 100644
--- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py
+++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py
@@ -1306,12 +1306,16 @@ def _stream_builder_model_map_cost(response: ModelResponse) -> float | None:
return None
+def _set_hidden_response_cost(response: ModelResponse, response_cost: float) -> None:
+ hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor
+ hidden_params["response_cost"] = response_cost
+
+
def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
response_cost: Final = _stream_builder_response_cost(response, logging_obj)
if response_cost is None:
return
- hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor
- hidden_params["response_cost"] = response_cost
+ _set_hidden_response_cost(response, response_cost)
def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
@@ -1320,8 +1324,10 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o
if isinstance(getattr(usage, "cost", None), (int, float)):
return
computed_cost: Final = logging_obj.response_cost_calculator(result=response)
- if isinstance(computed_cost, (int, float)) and computed_cost > 0:
- setattr(usage, "cost", computed_cost)
+ if not isinstance(computed_cost, (int, float)) or computed_cost <= 0:
+ return
+ setattr(usage, "cost", computed_cost)
+ _set_hidden_response_cost(response, float(computed_cost))
_NON_TEXT_DELTA_FIELDS: Final = (
diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py
index 164568451d4..de2cbf5f9e3 100644
--- a/litellm/llms/xai/cost_calculator.py
+++ b/litellm/llms/xai/cost_calculator.py
@@ -49,6 +49,19 @@ def _cost_reported_by_xai(usage: "Usage") -> float | None:
return float(reported_cost)
+def _split_reported_cost(model: str, usage: Usage, reported_cost: float) -> tuple[float, float]:
+ """xAI reports one total, so attribute it to input and output in the ratio our own token math gives."""
+ try:
+ priced_input, priced_output = _token_math_cost_per_token(model, usage)
+ except Exception: # noqa: BLE001 # generic_cost_per_token raises bare Exception for unmapped models
+ return 0.0, reported_cost
+ priced_total: Final = priced_input + priced_output
+ if priced_total <= 0:
+ return 0.0, reported_cost
+ input_cost: Final = reported_cost * priced_input / priced_total
+ return input_cost, reported_cost - input_cost
+
+
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
"""
Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens.
@@ -63,8 +76,11 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
"""
reported_cost: Final = _cost_reported_by_xai(usage)
if reported_cost is not None:
- return 0.0, reported_cost
+ return _split_reported_cost(model, usage, reported_cost)
+ return _token_math_cost_per_token(model, usage)
+
+def _token_math_cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
# XAI-specific completion cost: completion is billed as visible + reasoning
# tokens. Detect when the transformation layer already folded them so we
# don't double-count; fall back to raw xAI shape for callers that bypass
diff --git a/tests/integration/cost_calculation/cases/xai.py b/tests/integration/cost_calculation/cases/xai.py
index 2ee676fd4d0..3d0849e1311 100644
--- a/tests/integration/cost_calculation/cases/xai.py
+++ b/tests/integration/cost_calculation/cases/xai.py
@@ -19,7 +19,6 @@ XAI_GROK_4_7_STREAM: Final = CostTrackingTestCase(
name="xai-grok-4.7-stream",
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
model="xai/grok-4.7",
- known_bug="streamed xai responses bill the right total but the spend row splits it as input_cost 0 and output_cost = total: _stamp_streaming_usage_cost writes the stream total onto usage.cost and litellm/llms/xai/cost_calculator.py::cost_per_token reads usage.cost as an xAI reported price",
endpoint="/v1/chat/completions",
request={
"model": "$MODEL",
@@ -67,6 +66,7 @@ XAI_GROK_4_7_STREAM: Final = CostTrackingTestCase(
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
+ "cost_in_usd_ticks": 48600000,
"total_tokens": 2252,
"prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0},
"completion_tokens_details": {
@@ -83,7 +83,7 @@ XAI_GROK_4_7_STREAM: Final = CostTrackingTestCase(
),
),
expected=ExactExpected(
- spend=0.006152, input_cost=0.00368, output_cost=0.002472, prompt_tokens=1840, completion_tokens=412
+ spend=0.00486, input_cost=0.002907282, output_cost=0.001952718, prompt_tokens=1840, completion_tokens=412
),
)
@@ -148,6 +148,7 @@ XAI_GROK_4_7_TOOL_CALL: Final = CostTrackingTestCase(
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
+ "cost_in_usd_ticks": 61520000,
"total_tokens": 2252,
"prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0},
"completion_tokens_details": {
@@ -200,6 +201,7 @@ XAI_GROK_4_7_REASONING: Final = CostTrackingTestCase(
"usage": {
"prompt_tokens": 1840,
"completion_tokens": 412,
+ "cost_in_usd_ticks": 61520000,
"total_tokens": 2252,
"prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0},
"completion_tokens_details": {
@@ -224,6 +226,7 @@ XAI_GROK_4_7_CACHE_READ: Final = with_usage(
{
"prompt_tokens": 1840,
"completion_tokens": 412,
+ "cost_in_usd_ticks": 51920000,
"total_tokens": 2252,
"prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 640},
"completion_tokens_details": {
@@ -242,6 +245,7 @@ XAI_GROK_4_7_TIERED_INPUT_ABOVE_200K: Final = with_usage(
{
"prompt_tokens": 300000,
"completion_tokens": 412,
+ "cost_in_usd_ticks": 12049440000,
"total_tokens": 300412,
"prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0},
"completion_tokens_details": {
@@ -260,6 +264,7 @@ XAI_GROK_4_7_TIERED_CACHE_READ_ABOVE_200K: Final = with_usage(
{
"prompt_tokens": 300000,
"completion_tokens": 412,
+ "cost_in_usd_ticks": 6049440000,
"total_tokens": 300412,
"prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 200000},
"completion_tokens_details": {
@@ -278,6 +283,7 @@ XAI_GROK_4_7_TIERED_AT_200K_BOUNDARY: Final = with_usage(
{
"prompt_tokens": 200000,
"completion_tokens": 412,
+ "cost_in_usd_ticks": 8049440000,
"total_tokens": 200412,
"prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0},
"completion_tokens_details": {
diff --git a/tests/unit/llms/xai/test_xai_cost_calculator.py b/tests/unit/llms/xai/test_xai_cost_calculator.py
index 6065f1053e7..5219dba42ba 100644
--- a/tests/unit/llms/xai/test_xai_cost_calculator.py
+++ b/tests/unit/llms/xai/test_xai_cost_calculator.py
@@ -102,8 +102,8 @@ class TestXAICostCalculator:
def test_reported_cost_is_preferred_over_token_math(self):
"""The amount xAI reported, carried on usage.cost by the transformation, is billed.
- It lands entirely on completion cost because xAI does not split its total by
- direction, the same shape the perplexity calculator returns.
+ xAI does not split its total by direction, so the split follows the ratio our
+ own token math gives for the same usage while the sum stays xAI's number.
"""
usage = Usage(
prompt_tokens=100,
@@ -111,11 +111,15 @@ class TestXAICostCalculator:
total_tokens=300,
cost=0.0037756,
)
+ priced_prompt, priced_completion = cost_per_token(
+ model="grok-4.7", usage=Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
+ )
- prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage)
+ prompt_cost, completion_cost = cost_per_token(model="grok-4.7", usage=usage)
- assert prompt_cost == 0.0
- assert math.isclose(completion_cost, 0.0037756, rel_tol=1e-10)
+ assert math.isclose(prompt_cost + completion_cost, 0.0037756, rel_tol=1e-10)
+ assert math.isclose(prompt_cost / completion_cost, priced_prompt / priced_completion, rel_tol=1e-10)
+ assert not math.isclose(prompt_cost + completion_cost, priced_prompt + priced_completion, rel_tol=1e-3)
def test_reported_cost_suppresses_web_search_surcharge(self):
"""The reported total already covers server-side tool calls.
diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py
index 635ee98549e..7dc26fca1f1 100644
--- a/tests/unit/test_main.py
+++ b/tests/unit/test_main.py
@@ -4250,6 +4250,48 @@ def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch
assert response._hidden_params["response_cost"] == pytest.approx(usage_cost)
+def test_stream_chunk_builder_stamped_cost_is_not_re_read_as_xai_reported_cost():
+ import time as time_module
+
+ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
+
+ logging_obj: Final = LiteLLMLogging(
+ model="xai/grok-4.7",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=True,
+ call_type="completion",
+ start_time=time_module.time(),
+ litellm_call_id="stream-builder-xai-stamped-cost",
+ function_id="1",
+ )
+ logging_obj.model_call_details["custom_llm_provider"] = "xai"
+ logging_obj.optional_params = {}
+ usage_chunk: Final = _stream_builder_text_chunk("grok-4.7", "")
+ usage_chunk.usage = Usage(prompt_tokens=1840, completion_tokens=412, total_tokens=2252)
+ chunks: Final = [
+ _stream_builder_text_chunk("grok-4.7", "Hello.", finish_reason="stop"),
+ usage_chunk,
+ ]
+
+ response: Final = litellm.stream_chunk_builder(
+ chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj
+ )
+
+ assert response is not None
+ usage_cost: Final = getattr(response.usage, "cost", None)
+ assert usage_cost is not None
+ assert usage_cost > 0
+
+ spend_row_cost: Final = logging_obj._response_cost_calculator(result=response)
+ breakdown: Final = logging_obj.cost_breakdown
+ assert breakdown is not None
+ assert breakdown["input_cost"] > 0
+ assert breakdown["output_cost"] > 0
+ assert breakdown["input_cost"] + breakdown["output_cost"] == pytest.approx(usage_cost)
+ assert spend_row_cost == pytest.approx(usage_cost)
+ assert response._hidden_params.get("response_cost") == pytest.approx(usage_cost)
+
+
def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable():
import time as time_module
From 0fc11653e23782f778b82509f3f4682298493efe Mon Sep 17 00:00:00 2001
From: yujonglee
Date: Wed, 14 Oct 2026 17:07:20 -0700
Subject: [PATCH 060/140] test(rust-bridge): pin Messages and Responses
callback contracts on the native path (#46443)
* test(rust-bridge): pin Messages and Responses callback contracts on the native path
Each test runs the native route and asserts what the Python path does today, so a
known gap is a strict xfail that names its reason and the exception it fails with
Messages and Responses share one Surface and hook recorder in support/callbacks.py,
and each surface folder lists its own expected counts and gaps.
Body-level hooks from the function body (pre-request hook, websearch short-circuit,
Messages interceptors, planned and legacy agentic loops, streaming deployment hook)
are xfail until native runs them
Co-Authored-By: Claude Opus 5.5
* test(rust-bridge): give every route folder the same one-concern-per-file layout
Route data, the route Surface and the server fixture move into each route folder, and support/ keeps only
route-neutral helpers. test_inference.py splits into chat_completions/ and responses/, ocr/test_requests.py
splits into requests, failures and settings, and Messages and Responses get streaming and body hook files.
Messages request tests that duplicate inference-messages request.rs are removed, and drop_params read
from litellm.drop_params is now covered
Co-Authored-By: Claude Opus 5.5
* test(rust-bridge): expect the required-route decline error for unsupported Responses requests
Co-Authored-By: Yujong Lee
* test(rust-bridge): import route data from the route folders in the secrets unit tests
Co-Authored-By: Yujong Lee
* test(rust-bridge): import NoReturn in the OCR dispatch tests
Co-Authored-By: Yujong Lee
* test(rust-bridge): move route-neutral helpers out of support/callbacks.py
Co-Authored-By: Yujong Lee
* test(rust-bridge): cover deployment rejection and concurrent callbacks
* test(rust-bridge): assert native serving via the x-litellm-rust marker in the route folder tests
Co-Authored-By: Yujong Lee
---------
Co-authored-by: Claude Opus 5.5
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
tests/test_litellm_rust/AGENTS.md | 18 +
.../{support/cache.py => cache/support.py} | 4 +-
.../cache/test_python_cache.py | 11 +-
.../test_litellm_rust/cache/v2/test_cache.py | 8 +-
.../chat_completions/__init__.py | 0
.../chat_completions/conftest.py | 10 +
.../chat_completions/data.py | 28 +
.../chat_completions/test_callbacks.py | 65 ++
.../chat_completions/test_failures.py | 69 ++
.../chat_completions/test_lifecycle.py | 37 +
.../chat_completions/test_requests.py | 42 ++
.../chat_completions/test_settings.py | 82 +++
tests/test_litellm_rust/conftest.py | 4 +-
tests/test_litellm_rust/messages/conftest.py | 10 +
tests/test_litellm_rust/messages/data.py | 64 ++
.../messages/test_body_hooks.py | 137 ++++
.../messages/test_callbacks.py | 569 ++++++++++------
.../messages/test_request_shaping.py | 311 ---------
.../messages/test_requests.py | 150 +++++
.../messages/test_streaming.py | 73 ++
tests/test_litellm_rust/ocr/conftest.py | 19 +
tests/test_litellm_rust/ocr/data.py | 39 ++
tests/test_litellm_rust/ocr/test_callbacks.py | 181 +++--
tests/test_litellm_rust/ocr/test_cohere.py | 32 -
tests/test_litellm_rust/ocr/test_dispatch.py | 108 ++-
tests/test_litellm_rust/ocr/test_failures.py | 340 ++++++++++
.../test_litellm_rust/ocr/test_guardrails.py | 24 +-
tests/test_litellm_rust/ocr/test_lifecycle.py | 124 ++--
tests/test_litellm_rust/ocr/test_requests.py | 636 +-----------------
tests/test_litellm_rust/ocr/test_settings.py | 198 ++++++
tests/test_litellm_rust/responses/__init__.py | 0
tests/test_litellm_rust/responses/conftest.py | 10 +
tests/test_litellm_rust/responses/data.py | 75 +++
.../responses/test_body_hooks.py | 114 ++++
.../responses/test_callbacks.py | 408 +++++++++++
.../responses/test_failures.py | 83 +++
.../responses/test_lifecycle.py | 46 ++
.../responses/test_requests.py | 56 ++
.../responses/test_settings.py | 83 +++
.../responses/test_streaming.py | 58 ++
tests/test_litellm_rust/support/body_hooks.py | 39 ++
.../support/callback_recorder.py | 61 ++
tests/test_litellm_rust/support/gaps.py | 5 +
.../support/recording_server.py | 6 +
tests/test_litellm_rust/support/requests.py | 97 ---
tests/test_litellm_rust/support/results.py | 26 +
tests/test_litellm_rust/support/routing.py | 48 ++
tests/test_litellm_rust/support/served.py | 15 +-
tests/test_litellm_rust/test_inference.py | 403 -----------
.../unit/rust_bridge/messages/test_secrets.py | 2 +-
tests/unit/rust_bridge/ocr/test_secrets.py | 2 +-
51 files changed, 3157 insertions(+), 1873 deletions(-)
rename tests/test_litellm_rust/{support/cache.py => cache/support.py} (96%)
create mode 100644 tests/test_litellm_rust/chat_completions/__init__.py
create mode 100644 tests/test_litellm_rust/chat_completions/conftest.py
create mode 100644 tests/test_litellm_rust/chat_completions/data.py
create mode 100644 tests/test_litellm_rust/chat_completions/test_callbacks.py
create mode 100644 tests/test_litellm_rust/chat_completions/test_failures.py
create mode 100644 tests/test_litellm_rust/chat_completions/test_lifecycle.py
create mode 100644 tests/test_litellm_rust/chat_completions/test_requests.py
create mode 100644 tests/test_litellm_rust/chat_completions/test_settings.py
create mode 100644 tests/test_litellm_rust/messages/conftest.py
create mode 100644 tests/test_litellm_rust/messages/data.py
create mode 100644 tests/test_litellm_rust/messages/test_body_hooks.py
delete mode 100644 tests/test_litellm_rust/messages/test_request_shaping.py
create mode 100644 tests/test_litellm_rust/messages/test_requests.py
create mode 100644 tests/test_litellm_rust/messages/test_streaming.py
create mode 100644 tests/test_litellm_rust/ocr/conftest.py
create mode 100644 tests/test_litellm_rust/ocr/data.py
delete mode 100644 tests/test_litellm_rust/ocr/test_cohere.py
create mode 100644 tests/test_litellm_rust/ocr/test_failures.py
create mode 100644 tests/test_litellm_rust/ocr/test_settings.py
create mode 100644 tests/test_litellm_rust/responses/__init__.py
create mode 100644 tests/test_litellm_rust/responses/conftest.py
create mode 100644 tests/test_litellm_rust/responses/data.py
create mode 100644 tests/test_litellm_rust/responses/test_body_hooks.py
create mode 100644 tests/test_litellm_rust/responses/test_callbacks.py
create mode 100644 tests/test_litellm_rust/responses/test_failures.py
create mode 100644 tests/test_litellm_rust/responses/test_lifecycle.py
create mode 100644 tests/test_litellm_rust/responses/test_requests.py
create mode 100644 tests/test_litellm_rust/responses/test_settings.py
create mode 100644 tests/test_litellm_rust/responses/test_streaming.py
create mode 100644 tests/test_litellm_rust/support/body_hooks.py
create mode 100644 tests/test_litellm_rust/support/gaps.py
delete mode 100644 tests/test_litellm_rust/support/requests.py
create mode 100644 tests/test_litellm_rust/support/results.py
create mode 100644 tests/test_litellm_rust/support/routing.py
delete mode 100644 tests/test_litellm_rust/test_inference.py
diff --git a/tests/test_litellm_rust/AGENTS.md b/tests/test_litellm_rust/AGENTS.md
index d65ffd613aa..7e8cd67165f 100644
--- a/tests/test_litellm_rust/AGENTS.md
+++ b/tests/test_litellm_rust/AGENTS.md
@@ -1 +1,19 @@
This directory holds only the tests that cannot be written in the Rust code
+
+Each route folder (`ocr/`, `messages/`, `responses/`, `chat_completions/`) uses the same file names, one concern per file
+
+- `test_dispatch.py` which path served the call
+- `test_requests.py` Python inputs that reach the wire
+- `test_failures.py` errors raised to the caller and rejection before provider IO
+- `test_settings.py` globals, env, named credentials and their precedence
+- `test_streaming.py` stream relay, early close and stream-time logging
+- `test_callbacks.py` the logger hook contract
+- `test_body_hooks.py` handler body hooks such as interceptors and agentic loops
+- `test_lifecycle.py` cancellation, concurrency, context and nested calls
+
+A concern that exists for only one route gets its own file there, never a provider-named one
+
+- Route data and the route's `call` live in `/data.py`, and `/conftest.py` owns the `server` fixture
+- `support/` holds only route-neutral helpers. Never parametrize one test across routes, except in feature folders such as `cache/`
+- Call the public API inside `support.routing.native_route(...)`. Call `_native` directly only for a case the public API cannot reach
+- Check the native path with `support.routing.assert_served_natively`
diff --git a/tests/test_litellm_rust/support/cache.py b/tests/test_litellm_rust/cache/support.py
similarity index 96%
rename from tests/test_litellm_rust/support/cache.py
rename to tests/test_litellm_rust/cache/support.py
index 5deacc7869f..a30f6b68cc6 100644
--- a/tests/test_litellm_rust/support/cache.py
+++ b/tests/test_litellm_rust/cache/support.py
@@ -16,9 +16,9 @@ from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES
from litellm.rust_bridge.public_call import NativeCall
from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES
from litellm.types.utils import ModelResponse
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
+from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE
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
-from tests.test_litellm_rust.test_inference import RESPONSES_MODEL, RESPONSES_RESPONSE
def payload(value: object) -> object:
diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py
index 139859bdf55..9a0a874ddbf 100644
--- a/tests/test_litellm_rust/cache/test_python_cache.py
+++ b/tests/test_litellm_rust/cache/test_python_cache.py
@@ -21,14 +21,13 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
from litellm.rust_bridge.configuration import Rollout
-from litellm.rust_bridge.dispatch import call_hook
from litellm.rust_bridge.public_call import NativeCall
from litellm.types.caching import CachingSupportedCallTypes
-from tests.test_litellm_rust.support.cache import cache_key, collect, invoke, payload
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
+from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
+from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
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
-from tests.test_litellm_rust.test_inference import RESPONSES_RESPONSE
pytestmark = pytest.mark.requires_rust_extension
@@ -350,8 +349,6 @@ async def test_cache_lookup_uses_backend_request_callback_semantics(
route: Literal["chat", "messages", "responses"],
legacy: bool,
) -> None:
- from tests.test_litellm_rust.support.requests import request_body
-
class Rewrite(RecordingLogger):
temperature = 0.1
diff --git a/tests/test_litellm_rust/cache/v2/test_cache.py b/tests/test_litellm_rust/cache/v2/test_cache.py
index 5f930bd4ba7..0ba0a17007f 100644
--- a/tests/test_litellm_rust/cache/v2/test_cache.py
+++ b/tests/test_litellm_rust/cache/v2/test_cache.py
@@ -13,11 +13,11 @@ from litellm.caching.caching_handler import (
)
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.types.utils import ModelResponse
-from tests.test_litellm_rust.support.cache import cache_key, invoke, payload
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
+from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE, RESPONSES_TEXT
+from tests.test_litellm_rust.cache.support import cache_key, invoke, payload
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_MODEL, MESSAGES_RESPONSE
-from tests.test_litellm_rust.test_inference import RESPONSES_MODEL, RESPONSES_RESPONSE
pytestmark = pytest.mark.requires_rust_extension
@@ -153,7 +153,7 @@ async def test_v2_cache_does_not_force_native_responses_streaming(
assert isinstance(response, AsyncIterator)
chunks: Final = [chunk async for chunk in response]
assert chunks[-1].type == "response.completed"
- assert chunks[-1].response.output[0].content[0].text == "native response"
+ assert chunks[-1].response.output[0].content[0].text == RESPONSES_TEXT
@pytest.mark.asyncio
diff --git a/tests/test_litellm_rust/chat_completions/__init__.py b/tests/test_litellm_rust/chat_completions/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/test_litellm_rust/chat_completions/conftest.py b/tests/test_litellm_rust/chat_completions/conftest.py
new file mode 100644
index 00000000000..25c46c8ae42
--- /dev/null
+++ b/tests/test_litellm_rust/chat_completions/conftest.py
@@ -0,0 +1,10 @@
+import pytest
+
+from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+
+@pytest.fixture
+def server(recording_server: RecordingServer) -> RecordingServer:
+ recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
+ return recording_server
diff --git a/tests/test_litellm_rust/chat_completions/data.py b/tests/test_litellm_rust/chat_completions/data.py
new file mode 100644
index 00000000000..42f8e518f55
--- /dev/null
+++ b/tests/test_litellm_rust/chat_completions/data.py
@@ -0,0 +1,28 @@
+from typing import Final
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import call_natively
+
+
+def _error(kind: str) -> dict[str, object]:
+ return {"type": "error", "error": {"type": kind, "message": kind}}
+
+
+REQUEST: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32}
+BAD_REQUEST: Final = _error("invalid_request_error")
+RATE_LIMITED: Final = _error("rate_limit_error")
+
+
+async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object:
+ return await call_natively(
+ Route.CHAT_COMPLETIONS,
+ litellm.completion,
+ litellm.acompletion,
+ REQUEST,
+ server,
+ asynchronous=asynchronous,
+ **kwargs,
+ )
diff --git a/tests/test_litellm_rust/chat_completions/test_callbacks.py b/tests/test_litellm_rust/chat_completions/test_callbacks.py
new file mode 100644
index 00000000000..6743f9adbf0
--- /dev/null
+++ b/tests/test_litellm_rust/chat_completions/test_callbacks.py
@@ -0,0 +1,65 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.utils import CallTypes, ModelResponse
+from tests.test_litellm_rust.chat_completions.data import call
+from tests.test_litellm_rust.messages.data import MESSAGES_TEXT
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, request_body
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_returns_a_model_response_and_logs_success_once(server: RecordingServer, asynchronous: bool) -> None:
+ recorder: Final = RecordingLogger()
+
+ result: Final = await call(server, asynchronous=asynchronous, callbacks=[recorder], temperature=0.25)
+
+ assert_served_natively(server)
+ sent: Final = server.requests[0]
+ body: Final = _OBJECT.validate_python(sent.body)
+ assert sent.path == "/v1/messages"
+ assert body["temperature"] == 0.25
+ assert request_body(_OBJECT.validate_python(recorder.wait_for("log_pre_api_call")[0].kwargs)) == body
+ assert isinstance(result, ModelResponse)
+ assert result.choices[0].message.content == MESSAGES_TEXT
+ success: Final = await recorder.wait_for_async("async_log_success_event" if asynchronous else "log_success_event")
+ assert len(success) == 1
+ assert success[0].response is result
+
+
+@pytest.mark.asyncio
+async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None:
+ class Edit(CustomLogger):
+ def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
+ request_body(kwargs)["temperature"] = 0.75
+
+ await call(server, callbacks=[Edit()])
+
+ assert _OBJECT.validate_python(server.requests[0].body)["temperature"] == 0.75
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None:
+ class DropKeywords(CustomLogger):
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: CallTypes | None
+ ) -> dict[str, object]:
+ return {name: value for name, value in kwargs.items() if name not in ("temperature", "extra_headers")}
+
+ litellm.callbacks.append(DropKeywords())
+ await call(server, temperature=0.25, extra_headers={"x-deleted": "remove this header"})
+
+ assert_served_natively(server)
+ sent: Final = server.requests[0]
+ assert "temperature" not in _OBJECT.validate_python(sent.body)
+ assert sent.headers.get("x-deleted") is None
diff --git a/tests/test_litellm_rust/chat_completions/test_failures.py b/tests/test_litellm_rust/chat_completions/test_failures.py
new file mode 100644
index 00000000000..68a81096a7c
--- /dev/null
+++ b/tests/test_litellm_rust/chat_completions/test_failures.py
@@ -0,0 +1,69 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm import RateLimitError
+from tests.test_litellm_rust.chat_completions.data import RATE_LIMITED, call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+def failure_signal(asynchronous: bool) -> str:
+ return "async_log_failure_event" if asynchronous else "log_failure_event"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
+) -> None:
+ monkeypatch.setattr(litellm, "max_budget", 1.0)
+ monkeypatch.setattr(litellm, "_current_cost", 2.0)
+ monkeypatch.setattr(litellm, "ssl_verify", object())
+ recorder: Final = RecordingLogger()
+ server.expected_requests = 0
+
+ with pytest.raises(litellm.BudgetExceededError) as caught:
+ await call(server, asynchronous=asynchronous, callbacks=[recorder])
+
+ failure: Final = await recorder.wait_for_async(failure_signal(asynchronous))
+ assert len(failure) == 1
+ assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+ assert not server.requests
+ assert not any("success" in name for name in recorder.names)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_provider_failure_is_terminal_and_shared_with_callbacks(
+ server: RecordingServer, asynchronous: bool
+) -> None:
+ recorder: Final = RecordingLogger()
+ server.enqueue(ResponseSpec(body=RATE_LIMITED, status=429))
+
+ with pytest.raises(RateLimitError) as caught:
+ await call(server, asynchronous=asynchronous, callbacks=[recorder])
+
+ assert caught.value.status_code == 429
+ assert len(server.requests) == 1
+ failure: Final = await recorder.wait_for_async(failure_signal(asynchronous))
+ assert len(failure) == 1
+ assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+ assert not any("success" in name for name in recorder.names)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_empty_messages_fail_before_provider_io(server: RecordingServer, asynchronous: bool) -> None:
+ server.expected_requests = 0
+
+ with pytest.raises(Exception, match="chat completions requires at least one message"):
+ await call(server, asynchronous=asynchronous, messages=[])
+
+ assert not server.requests
diff --git a/tests/test_litellm_rust/chat_completions/test_lifecycle.py b/tests/test_litellm_rust/chat_completions/test_lifecycle.py
new file mode 100644
index 00000000000..f335d6fa509
--- /dev/null
+++ b/tests/test_litellm_rust/chat_completions/test_lifecycle.py
@@ -0,0 +1,37 @@
+import asyncio
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.rust_bridge import _native
+from litellm.rust_bridge.public_call import NativeCall
+from tests.test_litellm_rust.chat_completions.data import REQUEST
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+async def test_unstarted_call_has_no_provider_or_callback_effects(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorder: Final = RecordingLogger()
+ server.expected_requests = 0
+ monkeypatch.setattr(litellm, "ssl_verify", object())
+ monkeypatch.setattr(litellm, "max_budget", 1.0)
+ monkeypatch.setattr(litellm, "_current_cost", 2.0)
+ kwargs: Final = {
+ **REQUEST,
+ "api_key": "test-key",
+ "api_base": server.base_url,
+ "callbacks": [recorder],
+ }
+
+ pending: Final = _native.acompletion(NativeCall(args=(), kwargs=kwargs, base={}))
+ assert asyncio.iscoroutine(pending)
+ pending.close()
+
+ assert not server.requests
+ assert not recorder.events
diff --git a/tests/test_litellm_rust/chat_completions/test_requests.py b/tests/test_litellm_rust/chat_completions/test_requests.py
new file mode 100644
index 00000000000..3f2ba7c60e0
--- /dev/null
+++ b/tests/test_litellm_rust/chat_completions/test_requests.py
@@ -0,0 +1,42 @@
+import asyncio
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from litellm.types.utils import ModelResponse
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_TEXT
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively, native_route
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_positional_parameters_reach_the_wire(server: RecordingServer, asynchronous: bool) -> None:
+ temperature: Final = 0.35
+ arguments: Final = (
+ (MESSAGES_MODEL, list(MESSAGES), None, None, 12.0, temperature)
+ if asynchronous
+ else (MESSAGES_MODEL, list(MESSAGES), 12.0, temperature)
+ )
+ supplied: Final = {"base_url": server.base_url, "api_key": "test-key", "max_tokens": 32}
+
+ with native_route(Route.CHAT_COMPLETIONS):
+ result: Final = (
+ await litellm.acompletion(*arguments, **supplied)
+ if asynchronous
+ else await asyncio.to_thread(litellm.completion, *arguments, **supplied)
+ )
+
+ assert isinstance(result, ModelResponse)
+ assert result.choices[0].message.content == MESSAGES_TEXT
+ assert_served_natively(server)
+ body: Final = _OBJECT.validate_python(server.requests[0].body)
+ assert body["temperature"] == temperature
+ assert body["max_tokens"] == supplied["max_tokens"]
diff --git a/tests/test_litellm_rust/chat_completions/test_settings.py b/tests/test_litellm_rust/chat_completions/test_settings.py
new file mode 100644
index 00000000000..df64ab00421
--- /dev/null
+++ b/tests/test_litellm_rust/chat_completions/test_settings.py
@@ -0,0 +1,82 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.models.credentials import CredentialItem
+from litellm.types.utils import CallTypes, ModelResponse
+from tests.test_litellm_rust.chat_completions.data import call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("from_credentials", (False, True))
+async def test_resource_setup_uses_deployment_hook_arguments(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, from_credentials: bool
+) -> None:
+ invalid_settings: Final = {"ssl_verify": object()}
+ credential: Final = CredentialItem(
+ credential_name="resource-settings", credential_info={}, credential_values=invalid_settings
+ )
+ monkeypatch.setattr(litellm, "credential_list", [credential])
+
+ class Prepare(CustomLogger):
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: CallTypes | None
+ ) -> dict[str, object]:
+ return {
+ **kwargs,
+ **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings),
+ }
+
+ litellm.callbacks.append(Prepare())
+ recorder: Final = RecordingLogger()
+ server.expected_requests = 0
+
+ with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught:
+ await call(server, callbacks=[recorder])
+
+ failure: Final = await recorder.wait_for_async("async_log_failure_event")
+ assert len(failure) == 1
+ assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+ assert not server.requests
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty"))
+async def test_connection_settings_reach_the_provider(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, source: str
+) -> None:
+ key: Final = "selected-key"
+ explicit: Final = source in ("explicit", "base_url")
+ monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None))
+ monkeypatch.setattr(litellm, "anthropic_key", key if source == "provider" else ("unused" if explicit else None))
+ monkeypatch.setattr(
+ litellm,
+ "api_base",
+ None if source == "environment" else ("http://127.0.0.1:1" if explicit else server.base_url),
+ )
+ monkeypatch.setenv("ANTHROPIC_API_KEY", key if source == "environment" else "unused")
+ for name in ("ANTHROPIC_BASE_URL", "ANTHROPIC_API_BASE"):
+ monkeypatch.setenv(name, server.base_url if source == "environment" else "http://127.0.0.1:1")
+
+ result: Final = await call(
+ server,
+ asynchronous=asynchronous,
+ api_key=key if explicit else ("" if source == "empty" else None),
+ api_base=server.base_url if source == "explicit" else ("" if source == "empty" else None),
+ **({"base_url": server.base_url} if source == "base_url" else {}),
+ )
+
+ assert isinstance(result, ModelResponse)
+ assert_served_natively(server)
+ assert server.requests[0].headers["x-api-key"] == key
diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py
index 2e965fe2630..d9406af4c59 100644
--- a/tests/test_litellm_rust/conftest.py
+++ b/tests/test_litellm_rust/conftest.py
@@ -23,12 +23,12 @@ from tests.test_litellm_rust.support.recording_server import RecordingServer, re
@pytest_asyncio.fixture(autouse=True, loop_scope="function")
-async def isolate_ocr_test_state() -> AsyncIterator[None]:
+async def isolate_test_state() -> AsyncIterator[None]:
with ExitStack() as stack:
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")
+ executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-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))
diff --git a/tests/test_litellm_rust/messages/conftest.py b/tests/test_litellm_rust/messages/conftest.py
new file mode 100644
index 00000000000..25c46c8ae42
--- /dev/null
+++ b/tests/test_litellm_rust/messages/conftest.py
@@ -0,0 +1,10 @@
+import pytest
+
+from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+
+@pytest.fixture
+def server(recording_server: RecordingServer) -> RecordingServer:
+ recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
+ return recording_server
diff --git a/tests/test_litellm_rust/messages/data.py b/tests/test_litellm_rust/messages/data.py
new file mode 100644
index 00000000000..b7d039eb352
--- /dev/null
+++ b/tests/test_litellm_rust/messages/data.py
@@ -0,0 +1,64 @@
+from typing import Final
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from tests.test_litellm_rust.support.gaps import gap
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import call_natively
+
+MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5"
+MESSAGES: Final = ({"role": "user", "content": "Hello"},)
+MESSAGES_TEXT: Final = "Hello from native Messages"
+MESSAGES_RESPONSE: Final = {
+ "id": "msg_native",
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-sonnet-5",
+ "content": [{"type": "text", "text": MESSAGES_TEXT}],
+ "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": MESSAGES_TEXT}},
+ ),
+ ("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"}),
+)
+
+REQUEST: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 64}
+
+
+def _error(kind: str) -> dict[str, object]:
+ return {"type": "error", "error": {"type": kind, "message": kind}}
+
+
+BAD_REQUEST: Final = _error("invalid_request_error")
+RATE_LIMITED: Final = _error("rate_limit_error")
+
+
+async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object:
+ return await call_natively(
+ Route.MESSAGES,
+ litellm.anthropic.messages.create,
+ litellm.anthropic.messages.acreate,
+ REQUEST,
+ server,
+ asynchronous=asynchronous,
+ **kwargs,
+ )
+
+
+BODY_HOOKS: Final = gap("native Messages does not run the handler body hooks")
diff --git a/tests/test_litellm_rust/messages/test_body_hooks.py b/tests/test_litellm_rust/messages/test_body_hooks.py
new file mode 100644
index 00000000000..188cd58e00e
--- /dev/null
+++ b/tests/test_litellm_rust/messages/test_body_hooks.py
@@ -0,0 +1,137 @@
+from collections.abc import Iterator
+from datetime import datetime
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.anthropic.pass_through.messages import interceptors
+from litellm.llms.anthropic.pass_through.messages.interceptors.base import MessagesInterceptor
+from tests.test_litellm_rust.messages.data import BODY_HOOKS, MESSAGES_RESPONSE, REQUEST, call
+from tests.test_litellm_rust.support.callback_recorder import drain_logging
+from tests.test_litellm_rust.support.body_hooks import LegacyAgenticLoop, PlannedAgenticLoop, agentic_hooks
+from tests.test_litellm_rust.support.isolation import rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import response_id
+from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+ yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_pre_request_hook_rewrites_the_provider_request(server: RecordingServer) -> None:
+ class Rewrite(CustomLogger):
+ async def async_pre_request_hook(
+ self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object]
+ ) -> dict[str, object]:
+ return {**kwargs, "system": "body hook instruction"}
+
+ with rebound(litellm, "callbacks", [Rewrite()]):
+ await call(server)
+
+ assert_served_natively(server)
+ assert isinstance(server.requests[0].body, dict)
+ assert server.requests[0].body.get("system") == "body hook instruction"
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_websearch_short_circuit_answers_without_the_provider(server: RecordingServer) -> None:
+ class ShortCircuit(WebSearchInterceptionLogger):
+ async def try_short_circuit_search(
+ self,
+ model: str,
+ messages: list[dict[str, object]],
+ tools: list[dict[str, object]] | None,
+ custom_llm_provider: str | None,
+ kwargs: dict[str, object] | None = None,
+ ) -> object:
+ return {**MESSAGES_RESPONSE, "id": "short-circuited"}
+
+ with rebound(litellm, "callbacks", [ShortCircuit()]):
+ response: Final = await call(server)
+
+ assert response_id(response) == "short-circuited"
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_interceptor_replaces_the_provider_call(server: RecordingServer) -> None:
+ class Replace(MessagesInterceptor):
+ def can_handle(self, tools: list[dict[str, object]] | None, custom_llm_provider: str | None) -> bool:
+ return True
+
+ async def handle(self, **kwargs: object) -> object:
+ return {**MESSAGES_RESPONSE, "id": "intercepted"}
+
+ with rebound(interceptors, "_interceptors", [Replace()]):
+ response: Final = await call(server)
+
+ assert response_id(response) == "intercepted"
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_planned_agentic_loop_sends_the_followup_and_replaces_the_response(server: RecordingServer) -> None:
+ loop: Final = PlannedAgenticLoop()
+ server.enqueue(ResponseSpec(body=MESSAGES_RESPONSE))
+ server.default_response = ResponseSpec(body={**MESSAGES_RESPONSE, "id": "followup"})
+
+ with rebound(litellm, "callbacks", [loop]):
+ response: Final = await call(server)
+ await drain_logging()
+
+ assert_served_natively(server, count=2)
+ assert isinstance(server.requests[1].body, dict)
+ assert server.requests[1].body["messages"] == [{"role": "user", "content": "followup"}]
+ assert response_id(response) == "agentic-replaced"
+ assert agentic_hooks(loop) == (
+ "async_should_run_agentic_loop",
+ "async_build_agentic_loop_plan",
+ "async_should_run_agentic_loop",
+ "async_post_agentic_loop_response_hook",
+ )
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_legacy_agentic_loop_replaces_the_response(server: RecordingServer) -> None:
+ loop: Final = LegacyAgenticLoop()
+
+ with rebound(litellm, "callbacks", [loop]):
+ response: Final = await call(server)
+
+ assert_served_natively(server)
+ assert response_id(response) == "agentic-replaced"
+ assert agentic_hooks(loop) == ("async_should_run_agentic_loop", "async_run_agentic_loop")
+
+
+@pytest.mark.asyncio
+async def test_body_logging_setup_drops_a_logger_fn_from_a_supplied_logging_object(server: RecordingServer) -> None:
+ observed: Final[list[object]] = []
+ logging_obj: Final = Logging(
+ model=REQUEST["model"],
+ messages=[],
+ stream=False,
+ start_time=datetime.now(),
+ call_type="anthropic_messages",
+ litellm_call_id=CALL_ID,
+ function_id=CALL_ID,
+ )
+ logging_obj.update_environment_variables(litellm_params={"logger_fn": observed.append}, optional_params={})
+
+ await call(server, litellm_logging_obj=logging_obj)
+
+ assert_served_natively(server)
+ assert observed == []
diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py
index 6c317c3b9d9..e8d49453065 100644
--- a/tests/test_litellm_rust/messages/test_callbacks.py
+++ b/tests/test_litellm_rust/messages/test_callbacks.py
@@ -1,197 +1,138 @@
-from collections.abc import AsyncIterator, Iterator
+import asyncio
+from collections.abc import Iterator
+from contextvars import ContextVar
+from datetime import datetime
from typing import Final
import pytest
+from pydantic import TypeAdapter
import litellm
+from litellm.integrations.argilla import ArgillaLogger
from litellm.integrations.custom_logger import CustomLogger
-from litellm.rust_bridge import catalog
-from litellm.rust_bridge.catalog import Route, RouteRule
-from litellm.rust_bridge.configuration import Rollout
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from litellm.types.integrations.argilla import ArgillaItem
+from litellm.types.utils import ModelResponse, StandardLoggingPayload
+from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, call
+from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body
+from tests.test_litellm_rust.support.gaps import gap
from tests.test_litellm_rust.support.isolation import rebound
-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,
-)
-from tests.test_litellm_rust.support.served import assert_served_natively
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import response_id, with_id
+from tests.test_litellm_rust.support.routing import assert_served_natively
pytestmark = pytest.mark.requires_rust_extension
-STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS)
-
-
-@pytest.fixture(autouse=True)
-def opt_messages_into_rust() -> Iterator[None]:
- with rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), *catalog.RULES)):
- yield
+SYNC_HANDLERS_ON_ASYNC_CALLS: Final = gap("native async calls also run the sync success and failure handlers")
+_OBJECT: Final = TypeAdapter(dict[str, object])
+_HEADERS: Final = TypeAdapter(dict[str, str])
@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 server(server: RecordingServer) -> Iterator[RecordingServer]:
+ yield from relaxed(server)
@pytest.mark.asyncio
-async def test_native_messages_callbacks_see_the_provider_request_and_the_public_response(
- messages_server: RecordingServer,
+@pytest.mark.parametrize(
+ "hook,expected",
+ (
+ pytest.param("log_pre_api_call", 1),
+ pytest.param("log_post_api_call", 1),
+ pytest.param("async_log_success_event", 1),
+ pytest.param("async_logging_hook", 1),
+ pytest.param("redact_standard_logging_payload_from_model_call_details", 1, id="redact"),
+ pytest.param("async_pre_call_deployment_hook", 1),
+ pytest.param("async_post_call_success_deployment_hook", 1),
+ pytest.param("log_success_event", 0),
+ pytest.param("logging_hook", 0),
+ ),
+)
+async def test_success_hook_fires_the_expected_number_of_times(
+ server: RecordingServer, hook: str, expected: int
) -> None:
+ recorder: Final = HookRecorder()
+
+ with rebound(litellm, "callbacks", [recorder]):
+ response: Final = await call(server)
+ await recorder.wait_for_async("async_log_success_event")
+
+ assert_served_natively(server, response)
+ assert recorder.names.count(hook) == expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "hook,expected",
+ (
+ pytest.param("async_log_failure_event", 1),
+ pytest.param("async_post_call_failure_deployment_hook", 1),
+ pytest.param("log_model_group_rate_limit_error", 1),
+ pytest.param("log_failure_event", 0, marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+ ),
+)
+async def test_failure_hook_fires_the_expected_number_of_times(
+ server: RecordingServer, hook: str, expected: int
+) -> None:
+ server.default_response = ResponseSpec(body=RATE_LIMITED, status=429)
+ recorder: Final = HookRecorder()
+
+ with rebound(litellm, "callbacks", [recorder]):
+ with pytest.raises(litellm.RateLimitError):
+ await call(server, metadata={"model_group_size": 1})
+ await recorder.wait_for_async("async_log_failure_event")
+
+ assert_served_natively(server)
+ assert recorder.names.count(hook) == expected
+
+
+@pytest.mark.asyncio
+async def test_failure_events_carry_the_error_raised_to_the_caller(server: RecordingServer) -> None:
+ server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
recorder: Final = RecordingLogger()
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(messages_server, callbacks=[recorder], litellm_call_id="messages-success")
- )
+ with rebound(litellm, "callbacks", [recorder]):
+ with pytest.raises(litellm.BadRequestError) as raised:
+ await call(server)
+ events: Final = await recorder.wait_for_async("async_log_failure_event")
- assert_served_natively(messages_server, response)
- 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"
+ assert_served_natively(server)
+ assert isinstance(events[0].kwargs, dict)
+ assert events[0].kwargs["exception"] is raised.value
@pytest.mark.asyncio
-async def test_native_messages_pre_call_body_edit_reaches_the_provider(messages_server: RecordingServer) -> None:
+async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None:
+ recorder: Final = RecordingLogger()
+
+ response: Final = await call(server, callbacks=[recorder])
+
+ assert_served_natively(server, response)
+ pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs
+ assert isinstance(pre_call, dict)
+ assert request_body(pre_call) == server.requests[0].body
+
+
+@pytest.mark.asyncio
+async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None:
class Edit(CustomLogger):
- def log_pre_api_call(self, model, messages, kwargs):
+ def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
request_body(kwargs)["temperature"] = 0.25
- await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Edit()]))
+ response: Final = await call(server, callbacks=[Edit()])
- assert messages_server.requests[0].body["temperature"] == 0.25
+ assert_served_natively(server, response)
+ assert isinstance(server.requests[0].body, dict)
+ assert 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 len(messages_server.requests) == 1
- 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)
+async def test_pre_api_call_sees_the_shaped_optional_params(server: RecordingServer) -> None:
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]
+ await call(server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True)
- assert first + b"".join(rest) == sse_payload()
- assert_served_natively(messages_server, stream)
- 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, stream)
- 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, response)
- assert response["content"] == MESSAGES_RESPONSE["content"]
- assert len(recorder.wait_for("log_success_event")) == 1
-
-
-@pytest.mark.asyncio
-async def test_native_messages_pre_call_sees_the_shaped_optional_params(
- messages_server: RecordingServer,
-) -> None:
- recorder: Final = RecordingLogger()
-
- await litellm.anthropic.messages.acreate(
- **arguments(messages_server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True)
- )
-
- sent: Final = messages_server.requests[0].body
+ sent: Final = server.requests[0].body
+ assert isinstance(sent, dict)
assert not {"temperature", "top_k"} & sent.keys()
pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs
assert isinstance(pre_call, dict)
@@ -202,65 +143,287 @@ async def test_native_messages_pre_call_sees_the_shaped_optional_params(
@pytest.mark.asyncio
-async def test_native_messages_failing_pre_call_logger_does_not_fail_the_call(messages_server: RecordingServer) -> None:
- class Broken(CustomLogger):
- def log_pre_api_call(self, model, messages, kwargs):
- raise RuntimeError("logger exploded")
+async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None:
+ class Rewrite(HookRecorder):
+ async def async_logging_hook(
+ self, kwargs: dict[str, object], result: object, call_type: str
+ ) -> tuple[dict[str, object], object]:
+ return {**kwargs, "logging_hook_ran": True}, result
- response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Broken()]))
+ def redact_standard_logging_payload_from_model_call_details(
+ self, model_call_details: dict[str, object]
+ ) -> dict[str, object]:
+ return {**model_call_details, "redacted": True}
- assert_served_natively(messages_server, response)
- assert response["content"] == MESSAGES_RESPONSE["content"]
+ recorder: Final = Rewrite()
+
+ with rebound(litellm, "callbacks", [recorder]):
+ await call(server)
+ success: Final = await recorder.wait_for_async("async_log_success_event")
+
+ assert isinstance(success[0].kwargs, dict)
+ assert success[0].kwargs["logging_hook_ran"] is True
+ assert success[0].kwargs["redacted"] is True
@pytest.mark.asyncio
-async def test_native_messages_stream_success_log_carries_usage_rebuilt_from_the_relayed_events(
- messages_server: RecordingServer,
+async def test_pre_call_deployment_hooks_compose_the_wire_request_in_registration_order(
+ server: RecordingServer,
) -> None:
- messages_server.enqueue(STREAM)
+ class Rewrite(CustomLogger):
+ def __init__(self, marker: str) -> None:
+ super().__init__()
+ self.marker = marker
+
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: object
+ ) -> dict[str, object]:
+ headers: Final = _HEADERS.validate_python(kwargs["extra_headers"])
+ return {
+ **kwargs,
+ "extra_headers": {**headers, "x-deployment-hook": headers["x-deployment-hook"] + self.marker},
+ }
+
+ with rebound(litellm, "callbacks", [Rewrite("first"), Rewrite("-second")]):
+ response: Final = await call(server, extra_headers={"x-deployment-hook": "caller-", "x-kept": "unchanged"})
+
+ assert_served_natively(server, response)
+ assert server.requests[0].headers["x-deployment-hook"] == "caller-first-second"
+ assert server.requests[0].headers["x-kept"] == "unchanged"
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(server: RecordingServer) -> None:
+ rejection: Final = litellm.BadRequestError(
+ message="deployment rejected", model="test-model", llm_provider="anthropic"
+ )
recorder: Final = RecordingLogger()
- stream: Final = await litellm.anthropic.messages.acreate(
- **arguments(messages_server, stream=True, callbacks=[recorder])
- )
- assert isinstance(stream, AsyncIterator)
- async for _ in stream:
- pass
+ class Reject(CustomLogger):
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: object
+ ) -> dict[str, object]:
+ raise rejection
- success: Final = await recorder.wait_for_async("async_log_success_event")
- usage: Final = success[0].response.usage
- assert usage.completion_tokens == MESSAGES_EVENTS[4][1]["usage"]["output_tokens"]
- assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"]
- assert success[0].response.choices[0].message.content == "Hello from native Messages"
+ with rebound(litellm, "callbacks", [Reject(), recorder]):
+ with pytest.raises(litellm.BadRequestError, match="deployment rejected") as raised:
+ await call(server)
+ failures: Final = await recorder.wait_for_async("async_log_failure_event")
+
+ assert raised.value is rejection
+ assert server.requests == []
+ assert len(failures) == 1
+ assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value
+ assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names)
@pytest.mark.asyncio
-async def test_native_messages_does_not_restore_keywords_deleted_by_a_pre_call_hook(
- messages_server: RecordingServer,
-) -> None:
+async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None:
class DropKeywords(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: object
) -> dict[str, object]:
return {name: value for name, value in kwargs.items() if name not in ("system", "extra_headers")}
- call_arguments: Final = arguments(
- messages_server,
- system="remove this instruction",
- extra_headers={"x-deleted": "remove this header"},
- headers={"x-kept": "keep this header"},
- stream=False,
- )
with rebound(litellm, "callbacks", [DropKeywords()]):
- response: Final = await litellm.anthropic.messages.acreate(**call_arguments)
+ response: Final = await call(
+ server,
+ system="remove this instruction",
+ extra_headers={"x-deleted": "remove this header"},
+ headers={"x-kept": "keep this header"},
+ )
- assert_served_natively(messages_server, response)
- sent: Final = messages_server.requests[0]
- assert sent.body == {
- "model": MESSAGES_MODEL.removeprefix("anthropic/"),
- "messages": call_arguments["messages"],
- "max_tokens": call_arguments["max_tokens"],
- "stream": call_arguments["stream"],
- }
- assert sent.headers.get("x-deleted") is None
- assert sent.headers.get("x-kept") == "keep this header"
+ assert_served_natively(server, response)
+ sent: Final = server.requests[0]
+ assert isinstance(sent.body, dict)
+ assert "system" not in sent.body
+ assert "x-deleted" not in sent.headers
+ assert sent.headers["x-kept"] == "keep this header"
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_deployment_hook_replaces_the_response(server: RecordingServer) -> None:
+ recorder: Final = RecordingLogger()
+
+ class Replace(CustomLogger):
+ async def async_post_call_success_deployment_hook(
+ self, request_data: dict[str, object], response: object, call_type: object
+ ) -> object:
+ return with_id(response, "replaced")
+
+ with rebound(litellm, "callbacks", [Replace(), recorder]):
+ response: Final = await call(server)
+ success: Final = await recorder.wait_for_async("async_log_success_event")
+
+ assert_served_natively(server, response)
+ assert response_id(response) == "replaced"
+ assert len(success) == 1
+ assert isinstance(success[0].response, ModelResponse)
+ assert success[0].response.id == response_id(response)
+
+
+@pytest.mark.asyncio
+async def test_concurrent_calls_keep_callback_context_metadata_and_responses_isolated(server: RecordingServer) -> None:
+ context: Final = ContextVar("messages-callback-context", default="outside")
+ first_entered: Final = asyncio.Event()
+ second_entered: Final = asyncio.Event()
+ server.enqueue(ResponseSpec(body={**MESSAGES_RESPONSE, "id": "msg_first"}))
+ server.enqueue(ResponseSpec(body={**MESSAGES_RESPONSE, "id": "msg_second"}))
+
+ class Observe(RecordingLogger):
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: object
+ ) -> dict[str, object]:
+ arrived, other = (
+ (first_entered, second_entered)
+ if kwargs["litellm_call_id"] == "first"
+ else (second_entered, first_entered)
+ )
+ arrived.set()
+ await asyncio.wait_for(other.wait(), timeout=5)
+ return kwargs
+
+ def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
+ self._record("log_pre_api_call", {**kwargs, "caller_context": context.get()})
+
+ async def async_log_success_event(
+ self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+ ) -> None:
+ self._record("async_log_success_event", {**kwargs, "caller_context": context.get()}, response_obj)
+
+ recorder: Final = Observe()
+
+ async def invoke(marker: str) -> object:
+ token: Final = context.set(marker)
+ try:
+ return await call(server, litellm_call_id=marker, metadata={"marker": marker})
+ finally:
+ context.reset(token)
+
+ with rebound(litellm, "callbacks", [recorder]):
+ responses: Final = await asyncio.gather(invoke("first"), invoke("second"))
+ success: Final = await recorder.wait_for_async("async_log_success_event", count=2)
+
+ assert_served_natively(server, *responses)
+ expected: Final = {"first": responses[0], "second": responses[1]}
+ pre_call: Final = recorder.wait_for("log_pre_api_call", count=2)
+ assert len(pre_call) == len(success) == 2
+ assert {_OBJECT.validate_python(event.kwargs)["litellm_call_id"] for event in success} == expected.keys()
+ for event in (*pre_call, *success):
+ details: Final = _OBJECT.validate_python(event.kwargs)
+ call_id: Final = str(details["litellm_call_id"])
+ params: Final = _OBJECT.validate_python(details["litellm_params"])
+ assert details["caller_context"] == call_id
+ assert _OBJECT.validate_python(params["metadata"])["marker"] == call_id
+ for event in success:
+ logged_call_id: Final = str(_OBJECT.validate_python(event.kwargs)["litellm_call_id"])
+ assert isinstance(event.response, ModelResponse)
+ assert event.response.id == response_id(expected[logged_call_id])
+ assert context.get() == "outside"
+
+
+@pytest.mark.asyncio
+async def test_a_failing_logger_does_not_fail_the_call(server: RecordingServer) -> None:
+ class Broken(CustomLogger):
+ def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
+ raise RuntimeError("logger exploded")
+
+ response: Final = await call(server, callbacks=[Broken()])
+
+ assert_served_natively(server, response)
+ assert response_id(response) == MESSAGES_RESPONSE["id"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "event,asynchronous,dynamic,expected",
+ (
+ pytest.param("input", False, False, 1, id="global-input"),
+ pytest.param("success", True, False, 1, id="global-async-success"),
+ pytest.param("success", True, True, 1, id="request-async-success"),
+ pytest.param("success", False, False, 0, id="global-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+ pytest.param("success", False, True, 0, id="request-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+ pytest.param("failure", True, False, 1, id="global-async-failure"),
+ pytest.param("failure", False, False, 1, id="global-sync-failure"),
+ pytest.param("failure", False, True, 1, id="request-sync-failure"),
+ ),
+)
+async def test_function_callback_fires_the_expected_number_of_times(
+ server: RecordingServer, event: str, asynchronous: bool, dynamic: bool, expected: int
+) -> None:
+ recorder: Final = RecordingLogger()
+
+ def callback(kwargs: dict[str, object], *_: object) -> None:
+ recorder._record(event, kwargs)
+
+ async def async_callback(kwargs: dict[str, object], *_: object) -> None:
+ recorder._record(event, kwargs)
+
+ registered: Final = async_callback if asynchronous else callback
+ global_list: Final = f"_async_{event}_callback" if asynchronous else f"{event}_callback"
+ request_callbacks: Final = {f"{event}_callback": [registered]} if dynamic else {}
+ if event == "failure":
+ server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
+
+ with rebound(litellm, "callbacks", [recorder]), rebound(litellm, global_list, [] if dynamic else [registered]):
+ try:
+ await call(server, **request_callbacks)
+ except litellm.BadRequestError:
+ assert event == "failure"
+ await recorder.wait_for_async(f"async_log_{'failure' if event == 'failure' else 'success'}_event")
+
+ assert_served_natively(server)
+ assert recorder.names.count(event) == expected
+
+
+@pytest.mark.asyncio
+async def test_pre_call_rule_rejects_before_provider_io(server: RecordingServer) -> None:
+ checked: Final[list[str]] = []
+
+ def reject(text: str) -> bool:
+ checked.append(text)
+ return False
+
+ with rebound(litellm, "pre_call_rules", [reject]):
+ with pytest.raises(litellm.APIResponseValidationError):
+ await call(server)
+
+ assert checked == ["Hello"]
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+async def test_dataset_hook_returning_none_drops_the_record_from_the_success_sink(server: RecordingServer) -> None:
+ datasets: Final = RecordingLogger()
+
+ class Drop(CustomLogger):
+ async def async_dataset_hook(
+ self, logged_data: ArgillaItem, standard_logging_payload: StandardLoggingPayload | None
+ ) -> None:
+ datasets._record("async_dataset_hook", response=logged_data)
+
+ server.enqueue(ResponseSpec(body={"items": [{"id": "callback-dataset"}]}))
+ with rebound(litellm, "argilla_transformation_object", {"prompt": "messages", "answer": "response"}):
+ sink: Final = ArgillaLogger(
+ argilla_api_key="test-key",
+ argilla_base_url=server.base_url,
+ argilla_dataset_name="callback-dataset",
+ flush_interval=3600,
+ )
+ with rebound(litellm, "callbacks", [sink, Drop()]):
+ await call(server)
+ await datasets.wait_for_async("async_dataset_hook")
+ await sink.flush_queue()
+
+ assert [request.method for request in server.requests] == ["GET", "POST"]
+
+
+@pytest.mark.asyncio
+async def test_sync_call_returns_the_provider_message(server: RecordingServer) -> None:
+ recorder: Final = RecordingLogger()
+
+ response: Final = await call(server, asynchronous=False, callbacks=[recorder])
+
+ assert_served_natively(server, response)
+ assert response_id(response) == MESSAGES_RESPONSE["id"]
+ assert len(await recorder.wait_for_async("log_success_event")) == 1
diff --git a/tests/test_litellm_rust/messages/test_request_shaping.py b/tests/test_litellm_rust/messages/test_request_shaping.py
deleted file mode 100644
index 3ae17e0f614..00000000000
--- a/tests/test_litellm_rust/messages/test_request_shaping.py
+++ /dev/null
@@ -1,311 +0,0 @@
-"""The native Messages route shapes the wire request the way the Python handler does.
-
-Model capability expectations come from model_prices_and_context_window.json (Claude Sonnet 5 is an
-adaptive-thinking model without sampling params; Claude Haiku 4.5 is a legacy-thinking model), read at
-2026-09-24; the cost map is LiteLLM's own file.
-"""
-
-from collections.abc import Iterator
-from typing import Final
-
-import pytest
-
-import litellm
-from litellm.rust_bridge import catalog
-from litellm.rust_bridge.catalog import Route, RouteRule
-from litellm.rust_bridge.configuration import Rollout
-from tests.test_litellm_rust.support.isolation import rebound
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_RESPONSE
-from tests.test_litellm_rust.support.served import assert_served_natively
-
-pytestmark = pytest.mark.requires_rust_extension
-
-ADAPTIVE_MODEL: Final = "anthropic/claude-sonnet-5"
-LEGACY_THINKING_MODEL: Final = "anthropic/claude-haiku-4-5"
-
-
-@pytest.fixture(autouse=True)
-def opt_messages_into_rust() -> Iterator[None]:
- with rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), *catalog.RULES)):
- yield
-
-
-@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": ADAPTIVE_MODEL,
- "messages": [dict(message) for message in MESSAGES],
- "max_tokens": 8192,
- "api_key": "test-key",
- "api_base": server.base_url,
- **kwargs,
- }
-
-
-def sent(server: RecordingServer, response: object) -> tuple[dict[str, object], dict[str, str]]:
- assert_served_natively(server, response)
- request: Final = server.requests[0]
- assert isinstance(request.body, dict)
- return request.body, request.headers
-
-
-@pytest.mark.asyncio
-async def test_reasoning_effort_becomes_adaptive_thinking_and_effort_on_the_wire(
- messages_server: RecordingServer,
-) -> None:
- response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, reasoning_effort="high"))
-
- body, _ = sent(messages_server, response)
- assert "reasoning_effort" not in body
- assert body["thinking"] == {"type": "adaptive", "display": "summarized"}
- assert body["output_config"] == {"effort": "high"}
-
-
-@pytest.mark.asyncio
-async def test_claude_code_adaptive_payload_is_downgraded_to_a_capped_budget_for_a_legacy_model(
- messages_server: RecordingServer,
-) -> None:
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(
- messages_server,
- model=LEGACY_THINKING_MODEL,
- max_tokens=3000,
- thinking={"type": "adaptive"},
- output_config={"effort": "high"},
- temperature=0,
- )
- )
-
- body, _ = sent(messages_server, response)
- assert body["thinking"] == {"type": "enabled", "budget_tokens": 2999}
- assert "output_config" not in body
- assert "temperature" not in body
-
-
-@pytest.mark.asyncio
-async def test_removed_sampling_params_are_dropped_under_drop_params(messages_server: RecordingServer) -> None:
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(messages_server, temperature=0.2, top_p=0.9, top_k=5, drop_params=True)
- )
-
- body, _ = sent(messages_server, response)
- assert not {"temperature", "top_p", "top_k"} & body.keys()
-
-
-@pytest.mark.asyncio
-async def test_removed_sampling_params_are_rejected_without_drop_params(messages_server: RecordingServer) -> None:
- messages_server.expected_requests = 0
-
- with pytest.raises(litellm.BadRequestError, match="does not support top_k=5"):
- await litellm.anthropic.messages.acreate(**arguments(messages_server, top_k=5))
-
- assert messages_server.requests == []
-
-
-@pytest.mark.asyncio
-async def test_replayed_history_is_sanitized_before_it_reaches_the_provider(
- messages_server: RecordingServer,
-) -> None:
- history: Final = [
- {"role": "user", "content": "run it"},
- {
- "role": "assistant",
- "content": [
- {"type": "text", "text": ""},
- {
- "type": "tool_use",
- "id": "functions.Bash:0",
- "name": "Bash",
- "input": {},
- "provider_specific_fields": {"x": 1},
- },
- ],
- },
- {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]},
- ]
-
- response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, messages=history))
-
- body, _ = sent(messages_server, response)
- assert body["messages"] == [
- {"role": "user", "content": "run it"},
- {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]},
- {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]},
- ]
-
-
-@pytest.mark.asyncio
-async def test_feature_betas_merge_into_the_forwarded_beta_header(messages_server: RecordingServer) -> None:
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(
- messages_server,
- output_format={"type": "json_schema", "schema": {"type": "object"}},
- extra_headers={"anthropic-beta": "web-search-2025-03-05"},
- )
- )
-
- _, headers = sent(messages_server, response)
- assert headers["anthropic-beta"] == "structured-outputs-2025-11-13,web-search-2025-03-05"
-
-
-@pytest.mark.asyncio
-async def test_oauth_token_authenticates_as_a_bearer_with_the_oauth_beta(messages_server: RecordingServer) -> None:
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(messages_server, api_key="sk-ant-oat01-token")
- )
-
- _, headers = sent(messages_server, response)
- assert "x-api-key" not in headers
- assert headers["authorization"] == "Bearer sk-ant-oat01-token"
- assert headers["anthropic-beta"] == "oauth-2025-04-20"
- assert headers["anthropic-dangerous-direct-browser-access"] == "true"
-
-
-@pytest.mark.asyncio
-async def test_metadata_is_reduced_to_the_fields_anthropic_accepts(messages_server: RecordingServer) -> None:
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(messages_server, metadata={"user_id": "u-1", "trace_id": "internal"})
- )
-
- body, _ = sent(messages_server, response)
- assert body["metadata"] == {"user_id": "u-1"}
-
-
-@pytest.mark.asyncio
-async def test_additional_drop_params_remove_nested_fields_from_the_wire(messages_server: RecordingServer) -> None:
- tools: Final = [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}]
-
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(messages_server, tools=tools, additional_drop_params=["tools[*].input_examples"])
- )
-
- body, _ = sent(messages_server, response)
- assert body["tools"] == [{"name": "lookup", "input_schema": {"type": "object"}}]
-
-
-@pytest.mark.asyncio
-async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(messages_server: RecordingServer) -> None:
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(
- messages_server,
- provider_specific_header=[
- {"custom_llm_provider": "anthropic, azure_ai", "extra_headers": {"x-scoped": "yes"}},
- {"custom_llm_provider": "openai", "extra_headers": {"x-other": "no"}},
- ],
- )
- )
-
- _, headers = sent(messages_server, response)
- assert headers["x-scoped"] == "yes"
- assert "x-other" not in headers
-
-
-@pytest.mark.asyncio
-async def test_scoped_headers_override_extra_headers_which_override_forwarded_headers(
- messages_server: RecordingServer,
-) -> None:
- response: Final = await litellm.anthropic.messages.acreate(
- **arguments(
- messages_server,
- headers={"x-priority": "forwarded", "x-forwarded-only": "kept"},
- extra_headers={"x-priority": "extra", "x-extra-only": "kept"},
- provider_specific_header={"custom_llm_provider": "anthropic", "extra_headers": {"x-priority": "scoped"}},
- )
- )
-
- _, headers = sent(messages_server, response)
- assert {name: headers.get(name) for name in ("x-priority", "x-forwarded-only", "x-extra-only")} == {
- "x-priority": "scoped",
- "x-forwarded-only": "kept",
- "x-extra-only": "kept",
- }
-
-
-@pytest.mark.asyncio
-async def test_non_string_metadata_user_id_is_rejected_before_the_provider_call(
- messages_server: RecordingServer,
-) -> None:
- messages_server.expected_requests = 0
-
- with pytest.raises(litellm.BadRequestError, match=r"metadata\.user_id must be a string"):
- await litellm.anthropic.messages.acreate(**arguments(messages_server, metadata={"user_id": 123}))
-
- assert messages_server.requests == []
-
-
-@pytest.mark.asyncio
-async def test_native_messages_observes_runtime_capabilities_and_separate_caller_settings(
- messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
-) -> None:
- from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES
- from litellm.rust_bridge.public_call import NativeCall
-
- native: Final = NATIVE_AMESSAGES.load()
- assert native is not None
- model: Final = "claude-test-runtime-capabilities"
- messages_server.expected_requests = 2
- request: Final = NativeCall(
- args=(),
- kwargs={"temperature": 0.2, "drop_params": True},
- base={
- "model": model,
- "messages": MESSAGES,
- "max_tokens": 16,
- "stream": None,
- "api_key": "test-key",
- "api_base": messages_server.base_url,
- "custom_llm_provider": "anthropic",
- "temperature": 0.2,
- "drop_params": True,
- },
- )
- monkeypatch.setitem(
- litellm.model_cost,
- model,
- {
- "litellm_provider": "anthropic",
- "mode": "chat",
- "supports_sampling_params": True,
- },
- )
- first: Final = await native(request)
- monkeypatch.setitem(
- litellm.model_cost,
- model,
- {
- "litellm_provider": "anthropic",
- "mode": "chat",
- "supports_sampling_params": False,
- },
- )
- second: Final = await native(request)
-
- assert isinstance(first, dict)
- assert isinstance(second, dict)
- assert first["id"] == second["id"] == MESSAGES_RESPONSE["id"]
- assert len(messages_server.requests) == 2
- first_body: Final = messages_server.requests[0].body
- second_body: Final = messages_server.requests[1].body
- assert isinstance(first_body, dict)
- assert isinstance(second_body, dict)
- assert first_body["temperature"] == request.kwargs["temperature"]
- assert second_body == {name: value for name, value in first_body.items() if name != "temperature"}
-
-
-@pytest.mark.asyncio
-async def test_native_messages_reads_optional_positional_body_parameters(messages_server: RecordingServer) -> None:
- metadata: Final = {"user_id": "caller"}
- args: Final = (16, MESSAGES, "anthropic/claude-test", metadata, None, False, "Be brief", 0.25)
- kwargs: Final = {"api_key": "test-key", "api_base": messages_server.base_url}
- response: Final = await litellm.anthropic.messages.acreate(*args, **kwargs)
-
- body, _ = sent(messages_server, response)
- assert body["temperature"] == args[7]
- assert body["system"] == args[6]
- assert body["metadata"] == metadata
diff --git a/tests/test_litellm_rust/messages/test_requests.py b/tests/test_litellm_rust/messages/test_requests.py
new file mode 100644
index 00000000000..6388c0eb7db
--- /dev/null
+++ b/tests/test_litellm_rust/messages/test_requests.py
@@ -0,0 +1,150 @@
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_RESPONSE, call
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively, native_route
+
+pytestmark = pytest.mark.requires_rust_extension
+
+CAPABILITY_MODEL: Final = "claude-test-runtime-capabilities"
+
+
+def sent(server: RecordingServer) -> tuple[dict[str, object], dict[str, str]]:
+ assert_served_natively(server)
+ request: Final = server.requests[0]
+ assert isinstance(request.body, dict)
+ return request.body, request.headers
+
+
+def capabilities(monkeypatch: pytest.MonkeyPatch, *, sampling: bool) -> None:
+ monkeypatch.setitem(
+ litellm.model_cost,
+ CAPABILITY_MODEL,
+ {"litellm_provider": "anthropic", "mode": "chat", "supports_sampling_params": sampling},
+ )
+
+
+@pytest.mark.asyncio
+async def test_replayed_tool_ids_are_sanitized_before_they_reach_the_provider(server: RecordingServer) -> None:
+ history: Final = [
+ {"role": "user", "content": "run it"},
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "text", "text": ""},
+ {
+ "type": "tool_use",
+ "id": "functions.Bash:0",
+ "name": "Bash",
+ "input": {},
+ "provider_specific_fields": {"x": 1},
+ },
+ ],
+ },
+ {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]},
+ ]
+
+ await call(server, messages=history)
+
+ body, _ = sent(server)
+ assert body["messages"] == [
+ {"role": "user", "content": "run it"},
+ {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]},
+ {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]},
+ ]
+
+
+@pytest.mark.asyncio
+async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(server: RecordingServer) -> None:
+ await call(
+ server,
+ provider_specific_header=[
+ {"custom_llm_provider": "anthropic, azure_ai", "extra_headers": {"x-scoped": "yes"}},
+ {"custom_llm_provider": "openai", "extra_headers": {"x-other": "no"}},
+ ],
+ )
+
+ _, headers = sent(server)
+ assert headers["x-scoped"] == "yes"
+ assert "x-other" not in headers
+
+
+@pytest.mark.asyncio
+async def test_scoped_headers_override_extra_headers_which_override_forwarded_headers(server: RecordingServer) -> None:
+ await call(
+ server,
+ headers={"x-priority": "forwarded", "x-forwarded-only": "kept"},
+ extra_headers={"x-priority": "extra", "x-extra-only": "kept"},
+ provider_specific_header={"custom_llm_provider": "anthropic", "extra_headers": {"x-priority": "scoped"}},
+ )
+
+ _, headers = sent(server)
+ assert {name: headers.get(name) for name in ("x-priority", "x-forwarded-only", "x-extra-only")} == {
+ "x-priority": "scoped",
+ "x-forwarded-only": "kept",
+ "x-extra-only": "kept",
+ }
+
+
+@pytest.mark.asyncio
+async def test_model_capabilities_are_read_from_model_cost_at_call_time(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ server.expected_requests = 2
+ model: Final = f"anthropic/{CAPABILITY_MODEL}"
+
+ capabilities(monkeypatch, sampling=True)
+ await call(server, model=model, temperature=0.2, drop_params=True)
+ capabilities(monkeypatch, sampling=False)
+ await call(server, model=model, temperature=0.2, drop_params=True)
+
+ assert_served_natively(server, count=2)
+ first, second = (request.body for request in server.requests)
+ assert isinstance(first, dict)
+ assert first["temperature"] == 0.2
+ assert second == {name: value for name, value in first.items() if name != "temperature"}
+
+
+@pytest.mark.asyncio
+async def test_global_drop_params_drops_unsupported_params(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ capabilities(monkeypatch, sampling=False)
+ monkeypatch.setattr(litellm, "drop_params", True)
+
+ await call(server, model=f"anthropic/{CAPABILITY_MODEL}", top_k=5)
+
+ body, _ = sent(server)
+ assert "top_k" not in body
+
+
+@pytest.mark.asyncio
+async def test_unsupported_params_are_rejected_before_io_without_drop_params(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ capabilities(monkeypatch, sampling=False)
+ monkeypatch.setattr(litellm, "drop_params", False)
+ server.expected_requests = 0
+
+ with pytest.raises(litellm.BadRequestError, match="top_k"):
+ await call(server, model=f"anthropic/{CAPABILITY_MODEL}", top_k=5)
+
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+async def test_optional_positional_body_parameters_reach_the_wire(server: RecordingServer) -> None:
+ metadata: Final = {"user_id": "caller"}
+ args: Final = (16, list(MESSAGES), "anthropic/claude-test", metadata, None, False, "Be brief", 0.25)
+
+ with native_route(Route.MESSAGES):
+ response: Final = await litellm.anthropic_messages(*args, api_key="test-key", api_base=server.base_url)
+
+ assert isinstance(response, dict)
+ assert response["id"] == MESSAGES_RESPONSE["id"]
+ body, _ = sent(server)
+ assert (body["metadata"], body["system"], body["temperature"]) == (args[3], args[6], args[7])
diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py
new file mode 100644
index 00000000000..c700c093d1a
--- /dev/null
+++ b/tests/test_litellm_rust/messages/test_streaming.py
@@ -0,0 +1,73 @@
+from collections.abc import AsyncIterator, Iterator
+from typing import Final
+
+import pytest
+
+from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_RESPONSE, call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import consume
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS)
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+ yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+async def test_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk(
+ server: RecordingServer,
+) -> None:
+ server.enqueue(STREAM)
+ recorder: Final = RecordingLogger()
+
+ stream: Final = await call(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_served_natively(server)
+ assert first + b"".join(rest) == b"".join(STREAM.payloads())
+ success: Final = await recorder.wait_for_async("async_log_success_event")
+ assert len(success) == 1
+ assert isinstance(success[0].kwargs, dict)
+ assert success[0].kwargs["completion_start_time"] is not None
+ usage: Final = success[0].response.usage
+ assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"]
+ assert usage.completion_tokens == MESSAGES_RESPONSE["usage"]["output_tokens"]
+ assert "log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+async def test_sync_stream_relays_provider_events_and_logs_success_once(server: RecordingServer) -> None:
+ server.enqueue(STREAM)
+ recorder: Final = RecordingLogger()
+
+ stream: Final = await call(server, asynchronous=False, stream=True, callbacks=[recorder])
+ chunks: Final = await consume(stream)
+
+ assert_served_natively(server)
+ assert b"".join(chunks) == b"".join(STREAM.payloads())
+ assert len(await recorder.wait_for_async("async_log_success_event")) == 1
+
+
+@pytest.mark.asyncio
+async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> None:
+ server.enqueue(STREAM)
+ recorder: Final = RecordingLogger()
+
+ stream: Final = await call(server, stream=True, callbacks=[recorder])
+ assert isinstance(stream, AsyncIterator)
+ await anext(stream)
+ await stream.aclose()
+
+ assert len(await recorder.wait_for_async("async_log_success_event")) == 1
+ with pytest.raises(StopAsyncIteration):
+ await anext(stream)
diff --git a/tests/test_litellm_rust/ocr/conftest.py b/tests/test_litellm_rust/ocr/conftest.py
new file mode 100644
index 00000000000..6d3834516d7
--- /dev/null
+++ b/tests/test_litellm_rust/ocr/conftest.py
@@ -0,0 +1,19 @@
+from typing import Final
+
+import pytest
+
+from tests.test_litellm_rust.ocr.data import OCR_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+
+@pytest.fixture
+def server(recording_server: RecordingServer) -> RecordingServer:
+ recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
+ return recording_server
+
+
+@pytest.fixture(params=[False, True], ids=["python", "rust"])
+def ocr_backend(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> bool:
+ enabled: Final = bool(request.param)
+ monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
+ return enabled
diff --git a/tests/test_litellm_rust/ocr/data.py b/tests/test_litellm_rust/ocr/data.py
new file mode 100644
index 00000000000..0f039892245
--- /dev/null
+++ b/tests/test_litellm_rust/ocr/data.py
@@ -0,0 +1,39 @@
+from typing import Final
+
+import litellm
+from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+
+OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}
+OCR_MODEL: Final = "mistral/mistral-ocr-latest"
+OCR_TEXT: Final = "native OCR response"
+OCR_RESPONSE: Final = {
+ "pages": [{"index": 0, "markdown": OCR_TEXT, "images": [], "dimensions": None}],
+ "model": "mistral-ocr-latest",
+ "usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
+}
+
+
+def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
+ return {
+ "model": OCR_MODEL,
+ "document": dict(OCR_DOCUMENT),
+ "api_key": "test-key",
+ "api_base": server.base_url,
+ **kwargs,
+ }
+
+
+def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+ response: Final = litellm.ocr(**ocr_arguments(server, **kwargs))
+ if not isinstance(response, OCRResponse):
+ raise TypeError(f"Expected OCRResponse, got {type(response).__name__}")
+ return response
+
+
+async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
+ return await litellm.aocr(**ocr_arguments(server, **kwargs))
+
+
+async def call(server: RecordingServer, asynchronous: bool, **kwargs: object) -> OCRResponse:
+ return await call_aocr(server, **kwargs) if asynchronous else call_ocr(server, **kwargs)
diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py
index 9e3945f11b1..5f961288800 100644
--- a/tests/test_litellm_rust/ocr/test_callbacks.py
+++ b/tests/test_litellm_rust/ocr/test_callbacks.py
@@ -15,49 +15,38 @@ import litellm
from litellm.integrations.custom_guardrail import CustomGuardrail
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, rebound
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import (
+from tests.test_litellm_rust.ocr.data import (
OCR_DOCUMENT,
OCR_RESPONSE,
- call_native,
- call_native_aocr,
- call_native_ocr,
+ call,
+ call_aocr,
+ call_ocr,
+)
+from tests.test_litellm_rust.support.callback_recorder import (
+ RecordingLogger,
+ drain_logging,
request_body,
request_headers,
)
+from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
pytestmark = pytest.mark.requires_rust_extension
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
- recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
- return recording_server
-
-
-def call_native_ocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
- return call_native_ocr(server, callbacks=callbacks, **kwargs)
-
-
-async def call_native_aocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
- return await call_native_aocr(server, callbacks=callbacks, **kwargs)
-
-
-def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_server: RecordingServer) -> None:
+def test_native_ocr_pre_call_callback_receives_transformed_provider_request(server: RecordingServer) -> None:
observations: Final = []
class Observe(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
observations.append((model, copy.deepcopy(kwargs["additional_args"])))
- call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0])
+ call_ocr(server, callbacks=[Observe()], pages=[0])
assert len(observations) == 1
model, additional_args = observations[0]
assert model == "mistral-ocr-latest"
- assert additional_args["api_base"] == f"{ocr_server.base_url}/v1/ocr"
+ assert additional_args["api_base"] == f"{server.base_url}/v1/ocr"
assert additional_args["complete_input_dict"] == {
"model": "mistral-ocr-latest",
"document": OCR_DOCUMENT,
@@ -66,7 +55,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_
@pytest.mark.asyncio
-async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(ocr_server: RecordingServer) -> None:
+async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(server: RecordingServer) -> None:
class DropKeywords(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: object
@@ -74,17 +63,17 @@ async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(o
return {name: value for name, value in kwargs.items() if name not in ("pages", "extra_headers")}
with rebound(litellm, "callbacks", [DropKeywords()]):
- await call_native_aocr(ocr_server, pages=[0], extra_headers={"x-deleted": "remove this header"})
+ await call_aocr(server, pages=[0], extra_headers={"x-deleted": "remove this header"})
- assert len(ocr_server.requests) == 1
- sent: Final = ocr_server.requests[0]
+ assert len(server.requests) == 1
+ sent: Final = server.requests[0]
assert sent.body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
assert sent.headers.get("x-deleted") is None
@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"])
def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
- ocr_server: RecordingServer, raise_after_edit: bool
+ server: RecordingServer, raise_after_edit: bool
) -> None:
observed: Final = []
@@ -98,13 +87,13 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
def log_pre_api_call(self, model, messages, kwargs):
observed.append(copy.deepcopy(request_body(kwargs)))
- call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False)
+ call_ocr(server, callbacks=[Edit(), Observe()], include_image_base64=False)
assert observed[0]["include_image_base64"] is True
- assert ocr_server.requests[0].body["include_image_base64"] is True
+ assert server.requests[0].body["include_image_base64"] is True
-def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_server: RecordingServer) -> None:
+def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(server: RecordingServer) -> None:
observed: Final = []
class Edit(CustomLogger):
@@ -115,16 +104,16 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_
def log_pre_api_call(self, model, messages, kwargs):
observed.append(dict(request_headers(kwargs)))
- call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()])
+ call_ocr(server, callbacks=[Edit(), Observe()])
assert observed[0]["x-audit-tag"] == "reviewed"
- assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed"
+ assert server.requests[0].headers["x-audit-tag"] == "reviewed"
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references(
- ocr_server: RecordingServer, asynchronous: bool
+ server: RecordingServer, asynchronous: bool
) -> None:
original: Final = dict(OCR_DOCUMENT)
replacement_url: Final = "data:application/pdf;base64,ZGVm"
@@ -144,22 +133,20 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_
"model": "mistral/mistral-ocr-latest",
"document": original,
"api_key": "test-key",
- "api_base": ocr_server.base_url,
+ "api_base": server.base_url,
"callbacks": [Retain(), Edit()],
}
- response: Final = (
- await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
- )
+ response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
assert aliases == [True]
assert retained[0]["document_url"] == replacement_url
assert original["document_url"] == replacement_url
- assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url
+ assert server.requests[0].body["document"]["document_url"] == replacement_url
assert response.pages[0].markdown == "native OCR response"
def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
observed: Final = []
@@ -171,13 +158,13 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov
def log_pre_api_call(self, model, messages, kwargs):
observed.append(request_body(kwargs))
- call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()])
+ call_ocr(server, callbacks=[Rebind(), Observe()])
assert observed == [{"replacement": True}]
- assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
+ assert server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
-def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_server: RecordingServer) -> None:
+def test_native_ocr_callback_retained_body_observes_later_callback_mutation(server: RecordingServer) -> None:
queued: Final = []
class QueuePayload(CustomLogger):
@@ -188,12 +175,12 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_
def log_pre_api_call(self, model, messages, kwargs):
request_body(kwargs)["queued-edit"] = True
- call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()])
+ call_ocr(server, callbacks=[QueuePayload(), Edit()])
assert queued[0]["queued-edit"] is True
-def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(ocr_server: RecordingServer) -> None:
+def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(server: RecordingServer) -> None:
token: Final = object()
terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue()
finished: Final = threading.Event()
@@ -206,7 +193,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o
terminal_tokens.put(kwargs["test-token"])
finished.set()
- call_native_ocr_with_callbacks(ocr_server, [Stash()])
+ call_ocr(server, callbacks=[Stash()])
assert finished.wait(10)
assert terminal_tokens.get_nowait() is token
@@ -214,13 +201,13 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o
@pytest.mark.asyncio
async def test_native_aocr_success_callback_receives_call_id_metadata_and_response(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
recorder: Final = RecordingLogger()
- await call_native_aocr_with_callbacks(
- ocr_server,
- [recorder],
+ await call_aocr(
+ server,
+ callbacks=[recorder],
litellm_call_id="ocr-success",
metadata={"source": "callback-test"},
)
@@ -235,9 +222,9 @@ async def test_native_aocr_success_callback_receives_call_id_metadata_and_respon
@pytest.mark.asyncio
async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_response(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
- ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+ server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
observations: Final = []
class Observe(CustomLogger):
@@ -248,7 +235,7 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp
observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj))
with pytest.raises(litellm.InternalServerError):
- await call_native_aocr_with_callbacks(ocr_server, [Observe()])
+ await call_aocr(server, callbacks=[Observe()])
assert [observation[0] for observation in observations] == ["sync", "async"]
assert all(observation[1] == "aocr" for observation in observations)
@@ -257,12 +244,12 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp
@pytest.mark.asyncio
-async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_server: RecordingServer) -> None:
+async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(server: RecordingServer) -> None:
caller_loop: Final = asyncio.get_running_loop()
caller_thread: Final = threading.current_thread()
recorder: Final = RecordingLogger()
- await call_native_aocr_with_callbacks(ocr_server, [recorder])
+ await call_aocr(server, callbacks=[recorder])
events: Final = await recorder.wait_for_async("log_pre_api_call")
assert len(events) == 1
@@ -272,9 +259,9 @@ async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_
@pytest.mark.asyncio
async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_callback(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
- ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
+ server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
token: Final = object()
observed: Final = []
@@ -289,20 +276,20 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal
observed.append(("async", kwargs["request-token"]))
with pytest.raises(litellm.InternalServerError):
- await call_native_aocr_with_callbacks(ocr_server, [TrackInFlightRequest()])
+ await call_aocr(server, callbacks=[TrackInFlightRequest()])
assert [event for event, _ in observed] == ["sync", "async"]
assert all(observed_token is token for _, observed_token in observed)
def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registered_multiple_times(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
recorder: Final = RecordingLogger()
- call_native_ocr_with_callbacks(
- ocr_server,
- [recorder, recorder],
+ call_ocr(
+ server,
+ callbacks=[recorder, recorder],
success_callback=[recorder],
failure_callback=[recorder],
)
@@ -340,21 +327,21 @@ class ApplyEdits(CustomLogger):
@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]
+ server: RecordingServer, edits: dict[str, object]
) -> None:
- ocr_server.expected_requests = None
+ server.expected_requests = None
with isolated_callback_registries():
- call_native_ocr_with_callbacks(ocr_server, [ApplyEdits(MappingProxyType(edits))])
+ call_ocr(server, callbacks=[ApplyEdits(MappingProxyType(edits))])
- assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits}
+ assert 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:
+def test_native_ocr_sync_hooks_see_no_running_event_loop(server: RecordingServer, hook: str) -> None:
recorder: Final = RecordingLogger()
- call_native_ocr_with_callbacks(ocr_server, [recorder])
+ call_ocr(server, callbacks=[recorder])
[event] = recorder.wait_for(hook)
assert event.loop is None
@@ -364,7 +351,7 @@ def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingSe
@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
+ server: RecordingServer, asynchronous: bool
) -> None:
retained: Final = []
@@ -372,21 +359,21 @@ async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact(
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 call(server, asynchronous, callbacks=[Retain()])
await drain_logging()
gc.collect()
[(details, body, headers)] = retained
- assert body == ocr_server.requests[0].body
+ assert body == server.requests[0].body
assert headers
- assert all(ocr_server.requests[0].headers[name.lower()] == value for name, value in headers.items())
+ assert all(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
@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:
+async def test_native_ocr_success_callbacks_share_one_logging_payload(server: RecordingServer, family: str) -> None:
queued: Final = []
finished: Final = threading.Event()
@@ -421,7 +408,7 @@ async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server
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 call(server, family == "async", callbacks=[QueuePayload(), StripPayload()])
await drain_logging()
assert await asyncio.to_thread(finished.wait, 10)
@@ -430,7 +417,7 @@ async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server
@pytest.mark.asyncio
async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
token: Final = object()
observed: Final = []
@@ -455,7 +442,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f
litellm.callbacks.append(Block())
with pytest.raises(Blocked) as raised:
- await call_native_aocr(ocr_server)
+ await call_aocr(server)
await drain_logging()
assert observed == [("sync", token, raised.value), ("async", token, raised.value)]
@@ -464,7 +451,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f
@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(
- ocr_server: RecordingServer,
+ server: RecordingServer,
isolated_azure_auth: None,
asynchronous: bool,
) -> None:
@@ -497,27 +484,25 @@ async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context
"azure_ad_token_provider": provider,
"callbacks": [Edit()],
}
- response: Final = (
- await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
- )
+ response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
assert response.pages[0].markdown == "native OCR response"
assert observations == ["token", "pre_call"]
- assert ocr_server.requests[0].headers["authorization"] == "Bearer edited"
+ assert server.requests[0].headers["authorization"] == "Bearer edited"
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
- ocr_server: RecordingServer,
+ server: RecordingServer,
isolated_azure_auth: None,
asynchronous: bool,
) -> None:
- ocr_server.expected_requests = 2
+ server.expected_requests = 2
calls: Final = []
def provider() -> str:
calls.append("token")
- nested: Final = call_native_ocr(ocr_server)
+ nested: Final = call_ocr(server)
assert nested.pages[0].markdown == "native OCR response"
return "outer-token"
@@ -526,12 +511,10 @@ async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
"api_key": None,
"azure_ad_token_provider": provider,
}
- response: Final = (
- await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
- )
+ response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
assert response.pages[0].markdown == "native OCR response"
assert calls == ["token"]
- assert [request.headers["authorization"] for request in ocr_server.requests] == [
+ assert [request.headers["authorization"] for request in server.requests] == [
"Bearer test-key",
"Bearer outer-token",
]
@@ -539,10 +522,10 @@ async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
@pytest.mark.asyncio
async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error(
- ocr_server: RecordingServer,
+ server: RecordingServer,
isolated_azure_auth: None,
) -> None:
- ocr_server.expected_requests = 2
+ server.expected_requests = 2
async def request(token: str, fail: bool) -> object:
def provider() -> str:
@@ -550,8 +533,8 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error
raise ValueError(token)
return token
- return await call_native_aocr(
- ocr_server,
+ return await call_aocr(
+ server,
model="azure_ai/mistral-ocr-latest",
api_key=None,
azure_ad_token_provider=provider,
@@ -567,7 +550,7 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error
assert isinstance(responses[1], litellm.APIConnectionError)
assert "Failed to get Azure AD token: failed" in str(responses[1])
assert isinstance(responses[2], OCRResponse)
- assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [
+ assert sorted(request.headers["authorization"] for request in server.requests) == [
"Bearer first",
"Bearer second",
]
@@ -575,7 +558,7 @@ async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error
@pytest.mark.asyncio
async def test_native_azure_ocr_releases_token_provider_after_cancellation(
- ocr_server: RecordingServer,
+ server: RecordingServer,
isolated_azure_auth: None,
) -> None:
import gc
@@ -590,16 +573,16 @@ async def test_native_azure_ocr_releases_token_provider_after_cancellation(
async def invoke() -> weakref.ReferenceType[Provider]:
provider: Final = Provider()
reference: Final = weakref.ref(provider)
- ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1))
+ server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1))
task: Final = asyncio.create_task(
- call_native_aocr(
- ocr_server,
+ call_aocr(
+ server,
model="azure_ai/mistral-ocr-latest",
api_key=None,
azure_ad_token_provider=provider,
)
)
- await ocr_server.wait_for_requests(1)
+ await server.wait_for_requests(1)
assert reference() is provider
task.cancel()
with pytest.raises(asyncio.CancelledError):
diff --git a/tests/test_litellm_rust/ocr/test_cohere.py b/tests/test_litellm_rust/ocr/test_cohere.py
deleted file mode 100644
index 8474e971c6f..00000000000
--- a/tests/test_litellm_rust/ocr/test_cohere.py
+++ /dev/null
@@ -1,32 +0,0 @@
-from typing import Final
-
-import pytest
-
-import litellm
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-
-pytestmark = pytest.mark.requires_rust_extension
-MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")
-IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
-BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32}
-PAYLOAD: Final = {
- "pages": [
- {
- "index": 4,
- "markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]},
- },
- {"markdown": {"content": "page two"}},
- ],
- "meta": {"billed_units": {"pages": 3}},
-}
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model", MODELS)
-async def test_public_cohere_health_check(recording_server: RecordingServer, model: str) -> None:
- recording_server.enqueue(ResponseSpec(body=PAYLOAD))
- response: Final = await litellm.ahealth_check(
- model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr"
- )
- assert "error" not in response
- assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,")
diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py
index a29a45376d9..bb75782fb67 100644
--- a/tests/test_litellm_rust/ocr/test_dispatch.py
+++ b/tests/test_litellm_rust/ocr/test_dispatch.py
@@ -1,42 +1,36 @@
-from typing import Final
+from typing import Final, NoReturn
import pytest
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
+from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
-from tests.test_litellm_rust.support.served import assert_served_natively
+from tests.test_litellm_rust.support.routing import assert_served_natively
pytestmark = pytest.mark.requires_rust_extension
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
- recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
- return recording_server
-
-
@pytest.mark.parametrize("enabled", [False, True, None])
-def test_public_ocr_uses_native_route_independently_of_flag(ocr_server: RecordingServer, enabled: bool | None) -> None:
+def test_public_ocr_uses_native_route_independently_of_flag(server: RecordingServer, enabled: bool | None) -> None:
litellm.rust(enabled)
response: Final = litellm.ocr(
model=OCR_MODEL,
document=OCR_DOCUMENT,
api_key="test-key",
- api_base=ocr_server.base_url,
+ api_base=server.base_url,
)
assert isinstance(response, OCRResponse)
assert response.pages[0].markdown == "native OCR response"
- assert_served_natively(ocr_server, response)
+ assert_served_natively(server, response)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("caching", [None, False, True])
async def test_ocr_does_not_depend_on_chat_cache(
- ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None
) -> None:
from litellm.caching.caching import Cache
@@ -45,9 +39,93 @@ async def test_ocr_does_not_depend_on_chat_cache(
"model": OCR_MODEL,
"document": OCR_DOCUMENT,
"api_key": "test-key",
- "api_base": ocr_server.base_url,
+ "api_base": server.base_url,
"caching": caching,
}
response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
assert response.pages[0].markdown == "native OCR response"
- assert len(ocr_server.requests) == 1
+ assert len(server.requests) == 1
+
+
+@pytest.mark.parametrize("required", [False, True])
+@pytest.mark.parametrize("failure", ["invalid", "live", "schema"])
+def test_native_projection_errors_never_select_python(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, required: bool, failure: str
+) -> None:
+ import dataclasses
+ import ssl
+
+ from litellm.rust_bridge import runtime
+ from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
+ from litellm.rust_bridge.configuration import Rollout
+ from litellm.rust_bridge.host import settings
+ from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR
+ from litellm.rust_bridge.public_call import NativeCall
+
+ server.expected_requests = 0
+ snapshot: Final = dataclasses.replace(settings.http_settings(), user_agent=1)
+ if failure == "schema":
+ monkeypatch.setattr(settings, "http_settings", lambda: snapshot)
+ else:
+ monkeypatch.setattr(
+ litellm, "ssl_verify", ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if failure == "live" else object()
+ )
+ request: Final = NativeCall(
+ args=(),
+ kwargs={},
+ base={
+ "model": "mistral/mistral-ocr-latest",
+ "document": OCR_DOCUMENT,
+ "api_key": "test-key",
+ "api_base": server.base_url,
+ "timeout": None,
+ "custom_llm_provider": "mistral",
+ "extra_headers": None,
+ },
+ )
+
+ def python_fallback() -> NoReturn:
+ pytest.fail("projection failures must not select Python")
+
+ with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"):
+ runtime.run(
+ RouteContext(Route.OCR, provider="mistral"),
+ binding=NATIVE_OCR,
+ native=lambda native: native(request),
+ python=python_fallback,
+ rules=(RouteRule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),),
+ )
+ assert server.requests == []
+
+
+MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")
+
+
+IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
+
+
+BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32}
+
+
+PAYLOAD: Final = {
+ "pages": [
+ {
+ "index": 4,
+ "markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]},
+ },
+ {"markdown": {"content": "page two"}},
+ ],
+ "meta": {"billed_units": {"pages": 3}},
+}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model", MODELS)
+async def test_ocr_health_check_uses_the_native_route(recording_server: RecordingServer, model: str) -> None:
+ recording_server.enqueue(ResponseSpec(body=PAYLOAD))
+ response: Final = await litellm.ahealth_check(
+ model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr"
+ )
+ assert "error" not in response
+ assert_served_natively(recording_server)
+ assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,")
diff --git a/tests/test_litellm_rust/ocr/test_failures.py b/tests/test_litellm_rust/ocr/test_failures.py
new file mode 100644
index 00000000000..41373fb6e77
--- /dev/null
+++ b/tests/test_litellm_rust/ocr/test_failures.py
@@ -0,0 +1,340 @@
+import json
+from collections.abc import Callable
+from dataclasses import dataclass
+from io import BytesIO
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+from pydantic import JsonValue
+
+import litellm
+from tests.test_litellm_rust.ocr.data import (
+ OCR_RESPONSE,
+ call,
+ call_aocr,
+ call_ocr,
+)
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_ocr_contract_upstream_status(
+ server: RecordingServer,
+ ocr_backend: bool,
+ asynchronous: bool,
+) -> None:
+ upstream: Final = ResponseSpec(body={"detail": "invalid provider option"}, status=422)
+ server.enqueue(upstream)
+ arguments: Final = {
+ "model": "vertex_ai/mistral-ocr-latest",
+ "vertex_project": "test-project",
+ "vertex_location": "us-central1",
+ "num_retries": 0,
+ }
+ with pytest.raises(litellm.BadRequestError) as caught:
+ await call(server, asynchronous, **arguments)
+ assert caught.value.status_code == upstream.status
+ assert caught.value.response.status_code == upstream.status
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("preserved", ["body", "headers"])
+async def test_ocr_contract_provider_error_details(
+ server: RecordingServer,
+ ocr_backend: bool,
+ asynchronous: bool,
+ preserved: str,
+) -> None:
+ payload: Final = {"message": "rate limited"}
+ headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"}
+ server.enqueue(ResponseSpec(body=payload, status=429, headers=headers))
+ with pytest.raises(litellm.RateLimitError) as caught:
+ await call(server, asynchronous, num_retries=0)
+ response: Final = caught.value.response
+ assert isinstance(response, httpx.Response)
+ if preserved == "body":
+ assert response.content == json.dumps(payload).encode()
+ else:
+ for name, value in headers.items():
+ assert response.headers.get(name.lower()) == value
+ assert response.headers.get(name.upper()) == value
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_ocr_contract_invalid_response_format(
+ server: RecordingServer,
+ ocr_backend: bool,
+ asynchronous: bool,
+) -> None:
+ server.expected_requests = 0
+ with pytest.raises(litellm.UnsupportedParamsError) as caught:
+ await call(server, asynchronous, req_format="bogus", num_retries=0)
+ assert caught.value.status_code == 400
+ for value in ("req_format", "bogus", "native", "litellm"):
+ assert value in str(caught.value)
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+ "document,field",
+ [
+ ([], "document"),
+ ({"document_url": "https://example.com/a.pdf"}, "type"),
+ ({"type": "text"}, "type"),
+ ],
+)
+async def test_ocr_contract_malformed_document_is_actionable(
+ server: RecordingServer,
+ ocr_backend: bool,
+ asynchronous: bool,
+ document: JsonValue,
+ field: str,
+) -> None:
+ server.expected_requests = None
+ with pytest.raises(litellm.BadRequestError) as caught:
+ await call(server, asynchronous, document=document, num_retries=0)
+ assert caught.value.status_code == 400
+ assert field.lower() in str(caught.value).lower()
+ assert "NoneType: None" not in str(caught.value)
+ assert "indices must be" not in str(caught.value)
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("option,value,field", [("pages", [-1], "pages"), ("features", [1], "features")])
+async def test_ocr_contract_azure_invalid_options_are_bad_requests(
+ server: RecordingServer,
+ ocr_backend: bool,
+ asynchronous: bool,
+ option: str,
+ value: JsonValue,
+ field: str,
+) -> None:
+ server.expected_requests = 0
+ arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0}
+ with pytest.raises(litellm.BadRequestError) as caught:
+ await call(server, asynchronous, **arguments)
+ assert caught.value.status_code == 400
+ assert field in str(caught.value)
+ assert server.requests == []
+
+
+def test_native_ocr_maps_provider_400_with_public_provider_details(server: RecordingServer) -> None:
+ server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400))
+
+ with pytest.raises(litellm.BadRequestError) as caught:
+ call_ocr(server)
+
+ assert caught.value.status_code == 400
+ assert caught.value.model == "mistral-ocr-latest"
+ assert caught.value.llm_provider == "mistral"
+ assert "invalid OCR request" in str(caught.value)
+
+
+class TokenAbort(BaseException):
+ pass
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("failure", ["ordinary", "abort"], ids=["value-error", "base-exception"])
+async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request(
+ server: RecordingServer,
+ isolated_azure_auth: None,
+ asynchronous: bool,
+ failure: str,
+) -> None:
+ server.expected_requests = 0
+ calls: Final = []
+ recorder: Final = RecordingLogger()
+ original: Final = {"ordinary": ValueError("token unavailable"), "abort": TokenAbort("abort")}
+
+ def token_provider() -> object:
+ calls.append("token")
+ raise original[failure]
+
+ arguments: Final = {
+ "model": "azure_ai/mistral-ocr-latest",
+ "api_key": None,
+ "azure_ad_token_provider": token_provider,
+ "callbacks": [recorder],
+ }
+ expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError
+ with pytest.raises(expected) as caught:
+ await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
+ assert calls == ["token"]
+ assert server.requests == []
+ assert "log_pre_api_call" not in recorder.names
+ if failure == "ordinary":
+ assert "Failed to get Azure AD token: token unavailable" in str(caught.value)
+ assert isinstance(caught.value.__context__, RuntimeError)
+ assert caught.value.__context__.__cause__ is original[failure]
+ else:
+ assert caught.value is original[failure]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_native_file_preparation_preserves_reader_exception(server: RecordingServer, asynchronous: bool) -> None:
+ server.expected_requests = 0
+ failure: Final = RuntimeError("reader failed")
+
+ class Reader:
+ def read(self) -> bytes:
+ raise failure
+
+ document: Final = {"type": "file", "file": Reader()}
+ with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught:
+ await call_aocr(server, document=document) if asynchronous else call_ocr(server, document=document)
+ assert caught.value.__context__ is failure
+
+
+COHERE_IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
+
+
+FILE_SIZE_LIMIT: Final = 50 * 1024 * 1024
+
+
+class IntReader:
+ def read(self) -> int:
+ return 1
+
+
+def oversized_file(tmp_path: Path) -> Path:
+ path: Final = tmp_path / "large.pdf"
+ with path.open("wb") as stream:
+ stream.truncate(FILE_SIZE_LIMIT + 1)
+ return path
+
+
+def empty_token() -> str:
+ return ""
+
+
+def unused_token() -> str:
+ raise AssertionError("the token provider must not run")
+
+
+@dataclass(frozen=True, slots=True)
+class PublicFailure:
+ arguments: Callable[[Path], dict[str, object]]
+ error: type[Exception]
+ match: str
+ provider_requests: int = 0
+ response: ResponseSpec | None = None
+ cause: type[BaseException] | None = None
+
+
+PUBLIC_FAILURES: Final = {
+ "unknown-req-format": PublicFailure(
+ lambda _: {"req_format": "raw"}, litellm.BadRequestError, "Invalid `req_format`"
+ ),
+ "empty-file": PublicFailure(
+ lambda _: {"document": {"type": "file", "file": BytesIO(b"")}}, litellm.BadRequestError, "File is empty"
+ ),
+ "oversized-file": PublicFailure(
+ lambda tmp_path: {"document": {"type": "file", "file": oversized_file(tmp_path)}},
+ litellm.BadRequestError,
+ "exceeds the size limit",
+ ),
+ "missing-file": PublicFailure(
+ lambda tmp_path: {"document": {"type": "file", "file": tmp_path / "missing.pdf"}},
+ litellm.APIConnectionError,
+ "File not found",
+ cause=FileNotFoundError,
+ ),
+ "reader-returns-non-bytes": PublicFailure(
+ lambda _: {"document": {"type": "file", "file": IntReader()}},
+ litellm.APIConnectionError,
+ "bytes or str",
+ cause=TypeError,
+ ),
+ "cohere-non-image": PublicFailure(
+ lambda _: {"model": "cohere/parse-v5.0"}, litellm.BadRequestError, "only accepts `image_url`"
+ ),
+ "cohere-unknown-format": PublicFailure(
+ lambda _: {"model": "cohere/parse-v5.0", "document": COHERE_IMAGE, "output_format": "html"},
+ litellm.BadRequestError,
+ "output_format",
+ ),
+ "azure-missing-api-base": PublicFailure(
+ lambda _: {
+ "model": "azure_ai/mistral-ocr-latest",
+ "api_key": None,
+ "api_base": None,
+ "azure_ad_token_provider": unused_token,
+ },
+ litellm.APIConnectionError,
+ "Missing Azure AI API Base",
+ ),
+ "azure-empty-token": PublicFailure(
+ lambda _: {
+ "model": "azure_ai/mistral-ocr-latest",
+ "api_key": None,
+ "azure_ad_token": "static-token",
+ "azure_ad_token_provider": empty_token,
+ },
+ litellm.APIConnectionError,
+ "Missing Azure AI credentials",
+ ),
+ "upstream-500": PublicFailure(
+ lambda _: {},
+ litellm.InternalServerError,
+ "provider unavailable",
+ provider_requests=1,
+ response=ResponseSpec(body={"message": "provider unavailable"}, status=500),
+ ),
+ "invalid-provider-response": PublicFailure(
+ lambda _: {},
+ litellm.APIConnectionError,
+ "pages",
+ provider_requests=1,
+ response=ResponseSpec(body={"pages": "invalid"}),
+ ),
+ "response-over-limit": PublicFailure(
+ lambda _: {"max_response_bytes": len(json.dumps(OCR_RESPONSE).encode()) - 1},
+ litellm.APIConnectionError,
+ "OCR response exceeds the size limit",
+ provider_requests=1,
+ ),
+ "timeout": PublicFailure(
+ lambda _: {"timeout": 0.01},
+ litellm.Timeout,
+ "",
+ provider_requests=1,
+ response=ResponseSpec(body=OCR_RESPONSE, delay=0.2),
+ ),
+}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("failure", PUBLIC_FAILURES.values(), ids=PUBLIC_FAILURES.keys())
+async def test_native_failures_raise_the_public_exception_class(
+ server: RecordingServer,
+ isolated_azure_auth: None,
+ tmp_path: Path,
+ asynchronous: bool,
+ failure: PublicFailure,
+) -> None:
+ server.expected_requests = failure.provider_requests
+ if failure.response is not None:
+ server.enqueue(failure.response)
+
+ with pytest.raises(failure.error, match=failure.match) as caught:
+ await call(server, asynchronous, **failure.arguments(tmp_path))
+
+ assert len(server.requests) == failure.provider_requests
+ if failure.cause is not None:
+ assert isinstance(caught.value.__context__, failure.cause)
diff --git a/tests/test_litellm_rust/ocr/test_guardrails.py b/tests/test_litellm_rust/ocr/test_guardrails.py
index de4590ba202..cfe221d2994 100644
--- a/tests/test_litellm_rust/ocr/test_guardrails.py
+++ b/tests/test_litellm_rust/ocr/test_guardrails.py
@@ -8,19 +8,13 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
from litellm.types.utils import CallTypes
+from tests.test_litellm_rust.ocr.data import call_aocr
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_native, call_native_aocr
+from tests.test_litellm_rust.support.recording_server import RecordingServer
pytestmark = pytest.mark.requires_rust_extension
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
- recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
- return recording_server
-
-
class ReplaceOCRMarkdown(CustomGuardrail):
def __init__(self) -> None:
super().__init__(
@@ -36,7 +30,7 @@ class ReplaceOCRMarkdown(CustomGuardrail):
@pytest.mark.asyncio
async def test_native_aocr_post_call_content_filter_blocks_matching_markdown(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
guardrail: Final = ContentFilterGuardrail(
guardrail_name="block-native-ocr-markdown",
@@ -46,22 +40,22 @@ async def test_native_aocr_post_call_content_filter_blocks_matching_markdown(
litellm.callbacks.append(guardrail)
with pytest.raises(HTTPException, match="Content blocked") as blocked:
- await call_native_aocr(ocr_server, guardrails=[guardrail.guardrail_name])
+ await call_aocr(server, guardrails=[guardrail.guardrail_name])
assert blocked.value.status_code == 400
- assert len(ocr_server.requests) == 1
+ assert len(server.requests) == 1
@pytest.mark.asyncio
async def test_native_aocr_post_call_replacement_reaches_caller_and_success_callback(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
guardrail: Final = ReplaceOCRMarkdown()
recorder: Final = RecordingLogger()
litellm.callbacks.append(guardrail)
- response: Final = await call_native_aocr(
- ocr_server,
+ response: Final = await call_aocr(
+ server,
callbacks=[recorder],
guardrails=[guardrail.guardrail_name],
)
@@ -71,4 +65,4 @@ async def test_native_aocr_post_call_replacement_reaches_caller_and_success_call
assert response.pages[0].markdown == "Reviewed OCR"
assert len(success_events) == 1
assert success_events[0].response.pages[0].markdown == "Reviewed OCR"
- assert "guardrails" not in ocr_server.requests[0].body
+ assert "guardrails" not in server.requests[0].body
diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py
index 40bcb992b22..c262454a678 100644
--- a/tests/test_litellm_rust/ocr/test_lifecycle.py
+++ b/tests/test_litellm_rust/ocr/test_lifecycle.py
@@ -15,51 +15,45 @@ from litellm._logging import trace_id_var
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from tests.test_litellm_rust.ocr.data import OCR_RESPONSE, call_aocr, call_ocr
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_RESPONSE, call_aocr, call_ocr
pytestmark = pytest.mark.requires_rust_extension
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
- recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
- return recording_server
-
-
@pytest.mark.asyncio
-async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) -> None:
+async def test_proxy_metadata_remains_python_owned(server: RecordingServer) -> None:
from litellm.proxy._types import UserAPIKeyAuth
recorder: Final = RecordingLogger()
auth: Final = UserAPIKeyAuth(user_id="ocr-user")
response: Final = await call_aocr(
- ocr_server, callbacks=[recorder], metadata={"user_api_key_auth": auth}, shared_session=object()
+ server, callbacks=[recorder], metadata={"user_api_key_auth": auth}, shared_session=object()
)
events: Final = await recorder.wait_for_async("async_log_success_event")
assert response.pages[0].markdown == "native OCR response"
assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user"
- assert "metadata" not in ocr_server.requests[0].body
+ assert "metadata" not in server.requests[0].body
@pytest.mark.asyncio
async def test_request_level_custom_pricing_reaches_logging_params_and_bills_the_call(
- ocr_server: RecordingServer,
+ server: RecordingServer,
) -> None:
recorder: Final = RecordingLogger()
- response: Final = await call_aocr(ocr_server, callbacks=[recorder], ocr_cost_per_page=0.05)
+ response: Final = await call_aocr(server, callbacks=[recorder], ocr_cost_per_page=0.05)
events: Final = await recorder.wait_for_async("async_log_success_event")
assert response.usage_info is not None and response.usage_info.pages_processed == 1
assert events[0].kwargs["litellm_params"]["ocr_cost_per_page"] == 0.05
assert response._hidden_params["response_cost"] == pytest.approx(0.05)
- assert "ocr_cost_per_page" not in ocr_server.requests[0].body
+ assert "ocr_cost_per_page" not in server.requests[0].body
@pytest.mark.asyncio
-async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr_server: RecordingServer) -> None:
+async def test_response_replacement_finalized_before_dispatch_in_caller_task(server: RecordingServer) -> None:
caller: Final = asyncio.current_task()
context: Final = ContextVar("lifecycle-test", default="before")
observations: Final = []
@@ -76,11 +70,11 @@ async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr
return response.model_copy(update={"model": "replaced"})
litellm.callbacks.append(Replace())
- response: Final = await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="native-final")
+ response: Final = await call_aocr(server, callbacks=[recorder], litellm_call_id="native-final")
events: Final = await recorder.wait_for_async("async_log_success_event")
assert observations == [("pre", caller, "pre"), ("post", caller, "pre")]
assert context.get() == "pre"
- assert ocr_server.requests[0].body["pages"] == [2]
+ assert server.requests[0].body["pages"] == [2]
assert response.model == "replaced"
assert events[0].response is response
assert response._hidden_params["litellm_call_id"] == "native-final"
@@ -90,7 +84,7 @@ async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_metadata_failure_dispatches_only_failure_and_releases_logger(
- ocr_server: RecordingServer, asynchronous: bool
+ server: RecordingServer, asynchronous: bool
) -> None:
failure: Final = RuntimeError("metadata failed")
seen: Final = []
@@ -120,8 +114,8 @@ async def test_metadata_failure_dispatches_only_failure_and_releases_logger(
)
reference: Final = weakref.ref(logger)
with pytest.raises(RuntimeError) as caught:
- await call_aocr(ocr_server, litellm_logging_obj=logger) if asynchronous else call_ocr(
- ocr_server, litellm_logging_obj=logger
+ await call_aocr(server, litellm_logging_obj=logger) if asynchronous else call_ocr(
+ server, litellm_logging_obj=logger
)
assert caught.value is failure
failure.__traceback__ = None
@@ -132,12 +126,12 @@ async def test_metadata_failure_dispatches_only_failure_and_releases_logger(
gc.collect()
assert seen == ([("sync", failure), ("async", failure)] if asynchronous else [("sync", failure)])
assert reference() is None
- assert len(ocr_server.requests) == 1
+ assert len(server.requests) == 1
@pytest.mark.asyncio
-async def test_mapped_failure_identity_and_deployment_snapshot(ocr_server: RecordingServer) -> None:
- ocr_server.enqueue(ResponseSpec(body={"message": "unavailable"}, status=500))
+async def test_mapped_failure_identity_and_deployment_snapshot(server: RecordingServer) -> None:
+ server.enqueue(ResponseSpec(body={"message": "unavailable"}, status=500))
recorder: Final = RecordingLogger()
snapshots: Final = []
@@ -148,20 +142,20 @@ async def test_mapped_failure_identity_and_deployment_snapshot(ocr_server: Recor
litellm.callbacks.append(Observe())
with pytest.raises(litellm.InternalServerError) as caught:
- await call_aocr(ocr_server, callbacks=[recorder])
+ await call_aocr(server, callbacks=[recorder])
failures: Final = tuple(event for event in recorder.events if "failure" in event.name)
assert [event.name for event in failures] == ["log_failure_event", "async_log_failure_event"]
assert all(event.kwargs["exception"] is caught.value for event in failures)
assert caught.value.status_code == 500
assert snapshots[0] is not caught.value
assert snapshots[0].status_code == 418
- assert len(ocr_server.requests) == 1
+ assert len(server.requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("phase", ["pre", "http"])
async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch(
- ocr_server: RecordingServer, phase: str
+ server: RecordingServer, phase: str
) -> None:
entered: Final = asyncio.Event()
recorder: Final = RecordingLogger()
@@ -174,21 +168,21 @@ async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch(
litellm.callbacks.append(Pause())
if phase == "http":
- ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2))
+ server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2))
if phase == "pre":
- ocr_server.expected_requests = 0
+ server.expected_requests = 0
restored: Final = []
async def invoke():
trace_id_var.set("parent")
try:
- await call_aocr(ocr_server, callbacks=[recorder], litellm_trace_id="native-call")
+ await call_aocr(server, callbacks=[recorder], litellm_trace_id="native-call")
finally:
restored.append(trace_id_var.get())
task: Final = asyncio.create_task(invoke())
if phase == "http":
- await ocr_server.wait_for_requests(1)
+ await server.wait_for_requests(1)
else:
await asyncio.wait_for(entered.wait(), 5)
task.cancel()
@@ -201,9 +195,7 @@ async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch(
@pytest.mark.asyncio
@pytest.mark.parametrize("blocked", [False, True])
-async def test_deferred_logging_requires_release_and_runs_at_most_once(
- ocr_server: RecordingServer, blocked: bool
-) -> None:
+async def test_deferred_logging_requires_release_and_runs_at_most_once(server: RecordingServer, blocked: bool) -> None:
recorder: Final = RecordingLogger()
logger: Final = Logging(
model="mistral-ocr-latest",
@@ -216,7 +208,7 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(
dynamic_async_success_callbacks=[recorder],
)
logger._defer_async_logging = True
- response: Final = await call_aocr(ocr_server, litellm_logging_obj=logger)
+ response: Final = await call_aocr(server, litellm_logging_obj=logger)
await drain_logging()
assert "async_log_success_event" not in recorder.names
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked)
@@ -228,7 +220,7 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(
assert events[0].response is response
-def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: RecordingServer) -> None:
+def test_sync_success_uses_executor_and_copied_caller_context(server: RecordingServer) -> None:
context: Final = ContextVar("sync-lifecycle", default="missing")
context.set("caller")
thread: Final = threading.current_thread()
@@ -240,7 +232,7 @@ def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: Record
observations.append((threading.current_thread(), context.get(), response_obj))
finished.set()
- response: Final = call_ocr(ocr_server, callbacks=[Observe()])
+ response: Final = call_ocr(server, callbacks=[Observe()])
assert finished.wait(5)
assert observations[0][0] is not thread
assert observations[0][1] == "caller"
@@ -248,44 +240,44 @@ def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: Record
@pytest.mark.asyncio
-async def test_nested_native_calls_preserve_context_and_dispatch_each_outcome(ocr_server: RecordingServer) -> None:
- ocr_server.expected_requests = 2
+async def test_nested_native_calls_preserve_context_and_dispatch_each_outcome(server: RecordingServer) -> None:
+ server.expected_requests = 2
recorder: Final = RecordingLogger()
outcomes: Final = []
class Nested(CustomLogger):
async def async_pre_call_deployment_hook(self, kwargs, call_type):
if kwargs.get("litellm_call_id") == "outer":
- outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="inner"))
+ outcomes.append(await call_aocr(server, callbacks=[recorder], litellm_call_id="inner"))
litellm.callbacks.append(Nested())
- outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="outer"))
+ outcomes.append(await call_aocr(server, callbacks=[recorder], litellm_call_id="outer"))
events: Final = await recorder.wait_for_async("async_log_success_event", count=2)
assert [event.kwargs["litellm_call_id"] for event in events] == ["inner", "outer"]
assert events[0].response is outcomes[0]
assert events[1].response is outcomes[1]
- assert len(ocr_server.requests) == 2
+ assert len(server.requests) == 2
-def test_sync_pre_call_can_make_nested_native_request(ocr_server: RecordingServer) -> None:
- ocr_server.expected_requests = 2
+def test_sync_pre_call_can_make_nested_native_request(server: RecordingServer) -> None:
+ server.expected_requests = 2
observed: Final = []
class Nested(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
if kwargs["litellm_call_id"] == "outer-sync":
- observed.append(call_ocr(ocr_server, litellm_call_id="inner-sync"))
+ observed.append(call_ocr(server, litellm_call_id="inner-sync"))
- response: Final = call_ocr(ocr_server, callbacks=[Nested()], litellm_call_id="outer-sync")
+ response: Final = call_ocr(server, callbacks=[Nested()], litellm_call_id="outer-sync")
assert observed[0].pages[0].markdown == response.pages[0].markdown
- assert len(ocr_server.requests) == 2
+ assert len(server.requests) == 2
-def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_server: RecordingServer) -> None:
+def test_unstarted_native_coroutine_releases_input_without_reading_file(server: RecordingServer) -> None:
from litellm.ocr.dispatch import _public_request
from litellm.rust_bridge import _native
- ocr_server.expected_requests = 0
+ server.expected_requests = 0
effects: Final = []
class File:
@@ -308,7 +300,7 @@ def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_serv
@pytest.mark.asyncio
-async def test_file_read_happens_after_deployment_hook_in_caller_task(ocr_server: RecordingServer) -> None:
+async def test_file_read_happens_after_deployment_hook_in_caller_task(server: RecordingServer) -> None:
effects: Final = []
caller: Final = asyncio.current_task()
@@ -323,13 +315,13 @@ async def test_file_read_happens_after_deployment_hook_in_caller_task(ocr_server
effects.append(("hook", asyncio.current_task()))
litellm.callbacks.append(Deployment())
- await call_aocr(ocr_server, document={"type": "file", "file": File()})
+ await call_aocr(server, document={"type": "file", "file": File()})
assert effects == [("hook", caller), ("read", caller)]
@pytest.mark.asyncio
-async def test_failure_callbacks_continue_within_both_families(ocr_server: RecordingServer) -> None:
- ocr_server.enqueue(ResponseSpec(body={"message": "failed"}, status=500))
+async def test_failure_callbacks_continue_within_both_families(server: RecordingServer) -> None:
+ server.enqueue(ResponseSpec(body={"message": "failed"}, status=500))
observed: Final = []
class Broken(CustomLogger):
@@ -349,7 +341,7 @@ async def test_failure_callbacks_continue_within_both_families(ocr_server: Recor
observed.append(("following-async", kwargs["exception"]))
with pytest.raises(litellm.InternalServerError) as caught:
- await call_aocr(ocr_server, callbacks=[Broken(), Following()])
+ await call_aocr(server, callbacks=[Broken(), Following()])
assert [name for name, _ in observed] == ["broken-sync", "following-sync", "broken-async", "following-async"]
assert all(error is caught.value for _, error in observed)
@@ -392,17 +384,17 @@ async def test_cancelling_native_transport_closes_connection_before_return() ->
@pytest.mark.asyncio
-async def test_response_limit_is_enforced_at_the_public_boundary(ocr_server: RecordingServer) -> None:
+async def test_response_limit_is_enforced_at_the_public_boundary(server: RecordingServer) -> None:
limit: Final = len(json.dumps(OCR_RESPONSE).encode()) - 1
with pytest.raises(litellm.APIConnectionError, match="OCR response exceeds the size limit"):
- await call_aocr(ocr_server, max_response_bytes=limit)
- assert len(ocr_server.requests) == 1
+ await call_aocr(server, max_response_bytes=limit)
+ assert len(server.requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", [False, True])
async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_client_wrapper(
- ocr_server: RecordingServer,
+ server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
failure: bool,
created_loggers: list[Logging],
@@ -439,14 +431,14 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl
monkeypatch.setattr(litellm_logging, "executor", probe)
monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe)
if failure:
- ocr_server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500))
+ server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500))
trace_id_var.set("callback-free-parent")
arguments: Final = {"litellm_trace_id": "callback-free-call", "litellm_call_id": "callback-free-id"}
if failure:
with pytest.raises(litellm.InternalServerError):
- await call_aocr(ocr_server, **arguments)
+ await call_aocr(server, **arguments)
else:
- response: Final = await call_aocr(ocr_server, **arguments)
+ response: Final = await call_aocr(server, **arguments)
assert response.pages[0].markdown == "native OCR response"
assert response._hidden_params["litellm_call_id"] == "callback-free-id"
assert response._hidden_params["response_cost"] is not None
@@ -468,20 +460,18 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl
@pytest.mark.parametrize(
"registration", ["success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"]
)
-async def test_terminal_registration_added_during_http_is_observed(
- ocr_server: RecordingServer, registration: str
-) -> None:
+async def test_terminal_registration_added_during_http_is_observed(server: RecordingServer, registration: str) -> None:
failure: Final = "failure" in registration
observer: Final = RecordingLogger()
- ocr_server.enqueue(
+ server.enqueue(
ResponseSpec(
body={"message": "provider failed"} if failure else OCR_RESPONSE, status=500 if failure else 200, delay=0.1
)
)
task: Final = asyncio.create_task(
- asyncio.to_thread(call_ocr, ocr_server) if registration == "success_callback" else call_aocr(ocr_server)
+ asyncio.to_thread(call_ocr, server) if registration == "success_callback" else call_aocr(server)
)
- await ocr_server.wait_for_requests(1)
+ await server.wait_for_requests(1)
getattr(litellm, registration).append(observer)
if failure:
with pytest.raises(litellm.InternalServerError):
@@ -523,7 +513,7 @@ def created_loggers(monkeypatch: pytest.MonkeyPatch) -> list[Logging]:
@pytest.mark.asyncio
@pytest.mark.parametrize("consumer", ["logger_fn", "raw_global", "request_debug"])
async def test_explicit_logging_consumers_keep_request_and_response_payloads(
- ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging], consumer: str
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging], consumer: str
) -> None:
snapshots: Final[list[dict[str, object]]] = []
if consumer == "raw_global":
@@ -533,7 +523,7 @@ async def test_explicit_logging_consumers_keep_request_and_response_payloads(
"raw_global": {},
"request_debug": {"litellm_request_debug": True},
}[consumer]
- response: Final = await call_aocr(ocr_server, **arguments)
+ response: Final = await call_aocr(server, **arguments)
details: Final = created_loggers[0].model_call_details
assert details["additional_args"]["complete_input_dict"]["model"] == "mistral-ocr-latest"
assert json.loads(details["original_response"])["pages"][0]["markdown"] == response.pages[0].markdown
diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py
index 8c4f3a403d6..c7b7a6d4059 100644
--- a/tests/test_litellm_rust/ocr/test_requests.py
+++ b/tests/test_litellm_rust/ocr/test_requests.py
@@ -1,160 +1,36 @@
-import json
-from collections.abc import Callable
-from dataclasses import dataclass
from io import BytesIO
-from pathlib import Path
-from typing import Final, NoReturn
+from typing import Final
-import httpx
import pytest
-from pydantic import JsonValue
-import litellm
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import (
+from tests.test_litellm_rust.ocr.data import (
OCR_DOCUMENT,
OCR_RESPONSE,
- call_native,
- call_native_aocr,
- call_native_ocr,
+ call_aocr,
+ call_ocr,
)
-from tests.test_litellm_rust.support.served import assert_served_natively
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+from tests.test_litellm_rust.support.routing import assert_served_natively
pytestmark = pytest.mark.requires_rust_extension
-@pytest.fixture(params=[False, True], ids=["python", "rust"])
-def ocr_backend(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> bool:
- enabled: Final = bool(request.param)
- monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
- return enabled
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_ocr_contract_upstream_status(
- ocr_server: RecordingServer,
- ocr_backend: bool,
- asynchronous: bool,
-) -> None:
- upstream: Final = ResponseSpec(body={"detail": "invalid provider option"}, status=422)
- ocr_server.enqueue(upstream)
- arguments: Final = {
- "model": "vertex_ai/mistral-ocr-latest",
- "vertex_project": "test-project",
- "vertex_location": "us-central1",
- "num_retries": 0,
- }
- with pytest.raises(litellm.BadRequestError) as caught:
- await call_native(ocr_server, asynchronous, **arguments)
- assert caught.value.status_code == upstream.status
- assert caught.value.response.status_code == upstream.status
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("preserved", ["body", "headers"])
-async def test_ocr_contract_provider_error_details(
- ocr_server: RecordingServer,
- ocr_backend: bool,
- asynchronous: bool,
- preserved: str,
-) -> None:
- payload: Final = {"message": "rate limited"}
- headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"}
- ocr_server.enqueue(ResponseSpec(body=payload, status=429, headers=headers))
- with pytest.raises(litellm.RateLimitError) as caught:
- await call_native(ocr_server, asynchronous, num_retries=0)
- response: Final = caught.value.response
- assert isinstance(response, httpx.Response)
- if preserved == "body":
- assert response.content == json.dumps(payload).encode()
- else:
- for name, value in headers.items():
- assert response.headers.get(name.lower()) == value
- assert response.headers.get(name.upper()) == value
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_ocr_contract_invalid_response_format(
- ocr_server: RecordingServer,
- ocr_backend: bool,
- asynchronous: bool,
-) -> None:
- ocr_server.expected_requests = 0
- with pytest.raises(litellm.UnsupportedParamsError) as caught:
- await call_native(ocr_server, asynchronous, req_format="bogus", num_retries=0)
- assert caught.value.status_code == 400
- for value in ("req_format", "bogus", "native", "litellm"):
- assert value in str(caught.value)
- assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize(
- "document,field",
- [
- ([], "document"),
- ({"document_url": "https://example.com/a.pdf"}, "type"),
- ({"type": "text"}, "type"),
- ],
-)
-async def test_ocr_contract_malformed_document_is_actionable(
- ocr_server: RecordingServer,
- ocr_backend: bool,
- asynchronous: bool,
- document: JsonValue,
- field: str,
-) -> None:
- ocr_server.expected_requests = None
- with pytest.raises(litellm.BadRequestError) as caught:
- await call_native(ocr_server, asynchronous, document=document, num_retries=0)
- assert caught.value.status_code == 400
- assert field.lower() in str(caught.value).lower()
- assert "NoneType: None" not in str(caught.value)
- assert "indices must be" not in str(caught.value)
- assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("option,value,field", [("pages", [-1], "pages"), ("features", [1], "features")])
-async def test_ocr_contract_azure_invalid_options_are_bad_requests(
- ocr_server: RecordingServer,
- ocr_backend: bool,
- asynchronous: bool,
- option: str,
- value: JsonValue,
- field: str,
-) -> None:
- ocr_server.expected_requests = 0
- arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0}
- with pytest.raises(litellm.BadRequestError) as caught:
- await call_native(ocr_server, asynchronous, **arguments)
- assert caught.value.status_code == 400
- assert field in str(caught.value)
- assert ocr_server.requests == []
-
-
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/mistral-ocr-latest", "reducto/parse-v3"])
async def test_ocr_contract_native_format_supported(
- ocr_server: RecordingServer,
+ server: RecordingServer,
ocr_backend: bool,
asynchronous: bool,
model: str,
) -> None:
- ocr_server.expected_requests = None
+ server.expected_requests = None
payload: Final = (
{"result": {"chunks": [{"content": "native OCR response"}]}, "usage": {"num_pages": 1}}
if model.startswith("reducto/")
else OCR_RESPONSE
)
- ocr_server.default_response = ResponseSpec(body=payload)
+ server.default_response = ResponseSpec(body=payload)
arguments: Final = {
"model": model,
"req_format": "native",
@@ -163,43 +39,39 @@ async def test_ocr_contract_native_format_supported(
if model.startswith("reducto/")
else OCR_DOCUMENT,
}
- response: Final = (
- await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
- )
+ response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
assert response.pages[0].markdown == "native OCR response"
assert response.get_provider_native_response() == payload
- assert len(ocr_server.requests) == 1
+ assert len(server.requests) == 1
if ocr_backend:
- assert_served_natively(ocr_server, response)
+ assert_served_natively(server, response)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_ocr_contract_unknown_reducto_model_reaches_provider(
- ocr_server: RecordingServer,
+ server: RecordingServer,
ocr_backend: bool,
asynchronous: bool,
) -> None:
- ocr_server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}})
+ server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}})
arguments: Final = {
"model": "reducto/future-parse-model",
"document": {"type": "document_url", "document_url": "reducto://ready.pdf"},
"num_retries": 0,
}
- response: Final = (
- await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
- )
+ response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
assert response.model == "future-parse-model"
assert response.pages[0].markdown == "future model response"
- assert len(ocr_server.requests) == 1
- assert ocr_server.requests[0].path == "/parse"
- assert ocr_server.requests[0].body == {"input": "reducto://ready.pdf"}
+ assert len(server.requests) == 1
+ assert server.requests[0].path == "/parse"
+ assert server.requests[0].body == {"input": "reducto://ready.pdf"}
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token(
- ocr_server: RecordingServer, isolated_azure_auth: None, asynchronous: bool
+ server: RecordingServer, isolated_azure_auth: None, asynchronous: bool
) -> None:
calls: Final = []
@@ -212,478 +84,24 @@ async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token(
"api_key": None,
"azure_ad_token_provider": token_provider,
}
- response: Final = (
- await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
- )
+ response: Final = await call_aocr(server, **arguments) if asynchronous else call_ocr(server, **arguments)
assert calls == ["token"]
assert response.pages[0].markdown == "native OCR response"
- assert_served_natively(ocr_server, response)
- assert ocr_server.requests[0].headers["authorization"] == "Bearer callback-token"
+ assert_served_natively(server, response)
+ assert server.requests[0].headers["authorization"] == "Bearer callback-token"
-@pytest.fixture
-def ocr_server(recording_server: RecordingServer) -> RecordingServer:
- recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
- return recording_server
-
-
-def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: RecordingServer) -> None:
- ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400))
-
- with pytest.raises(litellm.BadRequestError) as caught:
- call_native_ocr(ocr_server)
-
- assert caught.value.status_code == 400
- assert caught.value.model == "mistral-ocr-latest"
- assert caught.value.llm_provider == "mistral"
- assert "invalid OCR request" in str(caught.value)
-
-
-def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(ocr_server: RecordingServer) -> None:
- response: Final = call_native_ocr(
- ocr_server,
+def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(server: RecordingServer) -> None:
+ response: Final = call_ocr(
+ server,
document={"type": "file", "file": BytesIO(b"abc"), "mime_type": "image/png"},
opaque_extension=object(),
)
assert response.pages[0].markdown == "native OCR response"
- assert_served_natively(ocr_server, response)
- assert ocr_server.requests[0].body == {
+ assert_served_natively(server, response)
+ assert server.requests[0].body == {
"model": "mistral-ocr-latest",
"document": {"type": "image_url", "image_url": "data:image/png;base64,YWJj"},
}
-
-
-class TokenAbort(BaseException):
- pass
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("failure", ["ordinary", "abort"], ids=["value-error", "base-exception"])
-async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callback_and_request(
- ocr_server: RecordingServer,
- isolated_azure_auth: None,
- asynchronous: bool,
- failure: str,
-) -> None:
- ocr_server.expected_requests = 0
- calls: Final = []
- recorder: Final = RecordingLogger()
- original: Final = {"ordinary": ValueError("token unavailable"), "abort": TokenAbort("abort")}
-
- def token_provider() -> object:
- calls.append("token")
- raise original[failure]
-
- arguments: Final = {
- "model": "azure_ai/mistral-ocr-latest",
- "api_key": None,
- "azure_ad_token_provider": token_provider,
- "callbacks": [recorder],
- }
- expected: Final = TokenAbort if failure == "abort" else litellm.APIConnectionError
- with pytest.raises(expected) as caught:
- await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
- assert calls == ["token"]
- assert ocr_server.requests == []
- assert "log_pre_api_call" not in recorder.names
- if failure == "ordinary":
- assert "Failed to get Azure AD token: token unavailable" in str(caught.value)
- assert isinstance(caught.value.__context__, RuntimeError)
- assert caught.value.__context__.__cause__ is original[failure]
- else:
- assert caught.value is original[failure]
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-@pytest.mark.parametrize(
- "override, expected_key",
- [
- ({}, "credential-key"),
- ({"api_key": "explicit-key"}, "explicit-key"),
- ({"api_key": None}, "environment-key"),
- ],
- ids=["inherit", "explicit", "explicit-none"],
-)
-async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments(
- ocr_server: RecordingServer,
- monkeypatch: pytest.MonkeyPatch,
- asynchronous: bool,
- override: dict[str, object],
- expected_key: str,
-) -> None:
- from litellm.models.credentials import CredentialItem
-
- pages: Final = [0]
- opaque: Final = object()
- monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
- monkeypatch.setattr(
- litellm,
- "credential_list",
- [
- CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}),
- CredentialItem(
- credential_name="ocr-test",
- credential_info={},
- credential_values={
- "api_key": "credential-key",
- "api_base": ocr_server.base_url,
- "pages": pages,
- "opaque": opaque,
- },
- ),
- CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}),
- ],
- )
-
- class Observer(RecordingLogger):
- def log_pre_api_call(self, model, messages, kwargs):
- super().log_pre_api_call(model, messages, kwargs)
- pages.append(2)
-
- arguments: Final = {
- "model": "mistral/mistral-ocr-latest",
- "document": OCR_DOCUMENT,
- "litellm_credential_name": "ocr-test",
- "callbacks": [Observer()],
- **override,
- }
- response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
- assert response.pages[0].markdown == "native OCR response"
- assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_key}"
- assert ocr_server.requests[0].body["pages"] == [0, 2]
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True])
-async def test_native_file_preparation_preserves_reader_exception(
- ocr_server: RecordingServer, asynchronous: bool
-) -> None:
- ocr_server.expected_requests = 0
- failure: Final = RuntimeError("reader failed")
-
- class Reader:
- def read(self) -> bytes:
- raise failure
-
- document: Final = {"type": "file", "file": Reader()}
- with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught:
- await call_native_aocr(ocr_server, document=document) if asynchronous else call_native_ocr(
- ocr_server, document=document
- )
- assert caught.value.__context__ is failure
-
-
-COHERE_IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}
-FILE_SIZE_LIMIT: Final = 50 * 1024 * 1024
-
-
-class IntReader:
- def read(self) -> int:
- return 1
-
-
-def oversized_file(tmp_path: Path) -> Path:
- path: Final = tmp_path / "large.pdf"
- with path.open("wb") as stream:
- stream.truncate(FILE_SIZE_LIMIT + 1)
- return path
-
-
-def empty_token() -> str:
- return ""
-
-
-def unused_token() -> str:
- raise AssertionError("the token provider must not run")
-
-
-@dataclass(frozen=True, slots=True)
-class PublicFailure:
- arguments: Callable[[Path], dict[str, object]]
- error: type[Exception]
- match: str
- provider_requests: int = 0
- response: ResponseSpec | None = None
- cause: type[BaseException] | None = None
-
-
-PUBLIC_FAILURES: Final = {
- "unknown-req-format": PublicFailure(
- lambda _: {"req_format": "raw"}, litellm.BadRequestError, "Invalid `req_format`"
- ),
- "empty-file": PublicFailure(
- lambda _: {"document": {"type": "file", "file": BytesIO(b"")}}, litellm.BadRequestError, "File is empty"
- ),
- "oversized-file": PublicFailure(
- lambda tmp_path: {"document": {"type": "file", "file": oversized_file(tmp_path)}},
- litellm.BadRequestError,
- "exceeds the size limit",
- ),
- "missing-file": PublicFailure(
- lambda tmp_path: {"document": {"type": "file", "file": tmp_path / "missing.pdf"}},
- litellm.APIConnectionError,
- "File not found",
- cause=FileNotFoundError,
- ),
- "reader-returns-non-bytes": PublicFailure(
- lambda _: {"document": {"type": "file", "file": IntReader()}},
- litellm.APIConnectionError,
- "bytes or str",
- cause=TypeError,
- ),
- "cohere-non-image": PublicFailure(
- lambda _: {"model": "cohere/parse-v5.0"}, litellm.BadRequestError, "only accepts `image_url`"
- ),
- "cohere-unknown-format": PublicFailure(
- lambda _: {"model": "cohere/parse-v5.0", "document": COHERE_IMAGE, "output_format": "html"},
- litellm.BadRequestError,
- "output_format",
- ),
- "azure-missing-api-base": PublicFailure(
- lambda _: {
- "model": "azure_ai/mistral-ocr-latest",
- "api_key": None,
- "api_base": None,
- "azure_ad_token_provider": unused_token,
- },
- litellm.APIConnectionError,
- "Missing Azure AI API Base",
- ),
- "azure-empty-token": PublicFailure(
- lambda _: {
- "model": "azure_ai/mistral-ocr-latest",
- "api_key": None,
- "azure_ad_token": "static-token",
- "azure_ad_token_provider": empty_token,
- },
- litellm.APIConnectionError,
- "Missing Azure AI credentials",
- ),
- "upstream-500": PublicFailure(
- lambda _: {},
- litellm.InternalServerError,
- "provider unavailable",
- provider_requests=1,
- response=ResponseSpec(body={"message": "provider unavailable"}, status=500),
- ),
- "invalid-provider-response": PublicFailure(
- lambda _: {},
- litellm.APIConnectionError,
- "pages",
- provider_requests=1,
- response=ResponseSpec(body={"pages": "invalid"}),
- ),
- "response-over-limit": PublicFailure(
- lambda _: {"max_response_bytes": len(json.dumps(OCR_RESPONSE).encode()) - 1},
- litellm.APIConnectionError,
- "OCR response exceeds the size limit",
- provider_requests=1,
- ),
- "timeout": PublicFailure(
- lambda _: {"timeout": 0.01},
- litellm.Timeout,
- "",
- provider_requests=1,
- response=ResponseSpec(body=OCR_RESPONSE, delay=0.2),
- ),
-}
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("failure", PUBLIC_FAILURES.values(), ids=PUBLIC_FAILURES.keys())
-async def test_native_failures_raise_the_public_exception_class(
- ocr_server: RecordingServer,
- isolated_azure_auth: None,
- tmp_path: Path,
- asynchronous: bool,
- failure: PublicFailure,
-) -> None:
- ocr_server.expected_requests = failure.provider_requests
- if failure.response is not None:
- ocr_server.enqueue(failure.response)
-
- with pytest.raises(failure.error, match=failure.match) as caught:
- await call_native(ocr_server, asynchronous, **failure.arguments(tmp_path))
-
- assert len(ocr_server.requests) == failure.provider_requests
- if failure.cause is not None:
- assert isinstance(caught.value.__context__, failure.cause)
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize(
- "name,value",
- [
- ("ssl_verify", object()),
- ("ssl_certificate", 1),
- ("ssl_certificate", ""),
- ("user_url_allowed_hosts", ["example.test", 1]),
- ],
-)
-async def test_native_settings_fail_before_provider_io(
- ocr_server: RecordingServer,
- monkeypatch: pytest.MonkeyPatch,
- asynchronous: bool,
- name: str,
- value: object,
-) -> None:
- ocr_server.expected_requests = 0
- monkeypatch.setattr(litellm, name, value)
- with pytest.raises(ValueError, match=r"http_settings|url_policy"):
- await call_native(ocr_server, asynchronous, num_retries=0)
- assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize(
- "name,value",
- [
- ("vertex_project", 1),
- ("vertex_location", ["region"]),
- ],
-)
-async def test_native_vertex_globals_fail_before_provider_io_only_for_vertex_calls(
- ocr_server: RecordingServer,
- monkeypatch: pytest.MonkeyPatch,
- asynchronous: bool,
- name: str,
- value: object,
-) -> None:
- ocr_server.expected_requests = 1
- monkeypatch.setattr(litellm, name, value)
- await call_native(ocr_server, asynchronous, num_retries=0)
- with pytest.raises(ValueError, match=r"provider_defaults"):
- await call_native(ocr_server, asynchronous, model="vertex_ai/mistral-ocr-latest", num_retries=0)
- assert len(ocr_server.requests) == 1
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_native_ssl_context_is_terminal_configuration(ocr_server: RecordingServer, asynchronous: bool) -> None:
- import ssl
-
- ocr_server.expected_requests = 0
- context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
- with pytest.raises(ValueError, match=r"request\.ssl_verify.*SSLContext"):
- await call_native(ocr_server, asynchronous, ssl_verify=context, num_retries=0)
- assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_native_settings_preserve_protocol_failures(
- ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
-) -> None:
- ocr_server.expected_requests = 0
- failure: Final = LookupError("settings truth test failed")
- cause: Final = RuntimeError("settings cause")
-
- class RaisesBool:
- def __bool__(self) -> bool:
- raise failure from cause
-
- monkeypatch.setattr(litellm, "force_ipv4", RaisesBool())
- with pytest.raises(LookupError) as caught:
- await call_native(ocr_server, asynchronous, num_retries=0)
- assert caught.value is failure
- assert caught.value.__cause__ is cause
- assert caught.value.__traceback__ is not None
- assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-async def test_native_settings_observe_mutation_between_calls(
- ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
-) -> None:
- monkeypatch.setattr(litellm, "force_ipv4", "yes")
- monkeypatch.setattr(litellm, "http2", 1)
- monkeypatch.setattr(litellm, "vertex_project", [])
- monkeypatch.setattr(litellm, "vertex_location", 0)
- monkeypatch.setattr(litellm, "user_url_allowed_hosts", "EXAMPLE.TEST.")
- response: Final = await call_native(ocr_server, asynchronous, num_retries=0)
- assert response.pages[0].markdown == "native OCR response"
- assert_served_natively(ocr_server, response)
- monkeypatch.setattr(litellm, "ssl_certificate", 1)
- with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate"):
- await call_native(ocr_server, asynchronous, num_retries=0)
- assert len(ocr_server.requests) == 1
-
-
-@pytest.mark.parametrize("required", [False, True])
-@pytest.mark.parametrize("failure", ["invalid", "live", "schema"])
-def test_native_projection_errors_never_select_python(
- ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, required: bool, failure: str
-) -> None:
- import dataclasses
- import ssl
-
- from litellm.rust_bridge import runtime
- from litellm.rust_bridge.host import settings
- from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule
- from litellm.rust_bridge.configuration import Rollout
- from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR
- from litellm.rust_bridge.public_call import NativeCall
-
- ocr_server.expected_requests = 0
- snapshot: Final = dataclasses.replace(settings.http_settings(), user_agent=1)
- if failure == "schema":
- monkeypatch.setattr(settings, "http_settings", lambda: snapshot)
- else:
- monkeypatch.setattr(
- litellm, "ssl_verify", ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if failure == "live" else object()
- )
- request: Final = NativeCall(
- args=(),
- kwargs={},
- base={
- "model": "mistral/mistral-ocr-latest",
- "document": OCR_DOCUMENT,
- "api_key": "test-key",
- "api_base": ocr_server.base_url,
- "timeout": None,
- "custom_llm_provider": "mistral",
- "extra_headers": None,
- },
- )
-
- def python_fallback() -> NoReturn:
- pytest.fail("projection failures must not select Python")
-
- with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"):
- runtime.run(
- RouteContext(Route.OCR, provider="mistral"),
- binding=NATIVE_OCR,
- native=lambda native: native(request),
- python=python_fallback,
- rules=(RouteRule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),),
- )
- assert ocr_server.requests == []
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
-@pytest.mark.parametrize("present", [False, True], ids=["missing", "invalid-pem"])
-async def test_native_client_certificate_is_validated_before_io(
- ocr_server: RecordingServer,
- monkeypatch: pytest.MonkeyPatch,
- tmp_path: Path,
- asynchronous: bool,
- present: bool,
-) -> None:
- ocr_server.expected_requests = 0
- certificate: Final = tmp_path / "client.pem"
- if present:
- certificate.write_text("invalid certificate")
- monkeypatch.setattr(litellm, "ssl_certificate", str(certificate))
- with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate.*PEM") as caught:
- await call_native(ocr_server, asynchronous, num_retries=0)
- assert str(certificate) not in str(caught.value)
- assert ocr_server.requests == []
diff --git a/tests/test_litellm_rust/ocr/test_settings.py b/tests/test_litellm_rust/ocr/test_settings.py
new file mode 100644
index 00000000000..7505468db32
--- /dev/null
+++ b/tests/test_litellm_rust/ocr/test_settings.py
@@ -0,0 +1,198 @@
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+import litellm
+from tests.test_litellm_rust.ocr.data import (
+ OCR_DOCUMENT,
+ call,
+)
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize(
+ "override, expected_key",
+ [
+ ({}, "credential-key"),
+ ({"api_key": "explicit-key"}, "explicit-key"),
+ ({"api_key": None}, "environment-key"),
+ ],
+ ids=["inherit", "explicit", "explicit-none"],
+)
+async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments(
+ server: RecordingServer,
+ monkeypatch: pytest.MonkeyPatch,
+ asynchronous: bool,
+ override: dict[str, object],
+ expected_key: str,
+) -> None:
+ from litellm.models.credentials import CredentialItem
+
+ pages: Final = [0]
+ opaque: Final = object()
+ monkeypatch.setenv("MISTRAL_API_KEY", "environment-key")
+ monkeypatch.setattr(
+ litellm,
+ "credential_list",
+ [
+ CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}),
+ CredentialItem(
+ credential_name="ocr-test",
+ credential_info={},
+ credential_values={
+ "api_key": "credential-key",
+ "api_base": server.base_url,
+ "pages": pages,
+ "opaque": opaque,
+ },
+ ),
+ CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}),
+ ],
+ )
+
+ class Observer(RecordingLogger):
+ def log_pre_api_call(self, model, messages, kwargs):
+ super().log_pre_api_call(model, messages, kwargs)
+ pages.append(2)
+
+ arguments: Final = {
+ "model": "mistral/mistral-ocr-latest",
+ "document": OCR_DOCUMENT,
+ "litellm_credential_name": "ocr-test",
+ "callbacks": [Observer()],
+ **override,
+ }
+ response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
+ assert response.pages[0].markdown == "native OCR response"
+ assert server.requests[0].headers["authorization"] == f"Bearer {expected_key}"
+ assert server.requests[0].body["pages"] == [0, 2]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+ "name,value",
+ [
+ ("ssl_verify", object()),
+ ("ssl_certificate", 1),
+ ("ssl_certificate", ""),
+ ("user_url_allowed_hosts", ["example.test", 1]),
+ ],
+)
+async def test_native_settings_fail_before_provider_io(
+ server: RecordingServer,
+ monkeypatch: pytest.MonkeyPatch,
+ asynchronous: bool,
+ name: str,
+ value: object,
+) -> None:
+ server.expected_requests = 0
+ monkeypatch.setattr(litellm, name, value)
+ with pytest.raises(ValueError, match=r"http_settings|url_policy"):
+ await call(server, asynchronous, num_retries=0)
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize(
+ "name,value",
+ [
+ ("vertex_project", 1),
+ ("vertex_location", ["region"]),
+ ],
+)
+async def test_native_vertex_globals_fail_before_provider_io_only_for_vertex_calls(
+ server: RecordingServer,
+ monkeypatch: pytest.MonkeyPatch,
+ asynchronous: bool,
+ name: str,
+ value: object,
+) -> None:
+ server.expected_requests = 1
+ monkeypatch.setattr(litellm, name, value)
+ await call(server, asynchronous, num_retries=0)
+ with pytest.raises(ValueError, match=r"provider_defaults"):
+ await call(server, asynchronous, model="vertex_ai/mistral-ocr-latest", num_retries=0)
+ assert len(server.requests) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_native_ssl_context_is_terminal_configuration(server: RecordingServer, asynchronous: bool) -> None:
+ import ssl
+
+ server.expected_requests = 0
+ context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+ with pytest.raises(ValueError, match=r"request\.ssl_verify.*SSLContext"):
+ await call(server, asynchronous, ssl_verify=context, num_retries=0)
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_native_settings_preserve_protocol_failures(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
+) -> None:
+ server.expected_requests = 0
+ failure: Final = LookupError("settings truth test failed")
+ cause: Final = RuntimeError("settings cause")
+
+ class RaisesBool:
+ def __bool__(self) -> bool:
+ raise failure from cause
+
+ monkeypatch.setattr(litellm, "force_ipv4", RaisesBool())
+ with pytest.raises(LookupError) as caught:
+ await call(server, asynchronous, num_retries=0)
+ assert caught.value is failure
+ assert caught.value.__cause__ is cause
+ assert caught.value.__traceback__ is not None
+ assert server.requests == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+async def test_native_settings_observe_mutation_between_calls(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
+) -> None:
+ monkeypatch.setattr(litellm, "force_ipv4", "yes")
+ monkeypatch.setattr(litellm, "http2", 1)
+ monkeypatch.setattr(litellm, "vertex_project", [])
+ monkeypatch.setattr(litellm, "vertex_location", 0)
+ monkeypatch.setattr(litellm, "user_url_allowed_hosts", "EXAMPLE.TEST.")
+ response: Final = await call(server, asynchronous, num_retries=0)
+ assert response.pages[0].markdown == "native OCR response"
+ assert_served_natively(server)
+ monkeypatch.setattr(litellm, "ssl_certificate", 1)
+ with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate"):
+ await call(server, asynchronous, num_retries=0)
+ assert len(server.requests) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("present", [False, True], ids=["missing", "invalid-pem"])
+async def test_native_client_certificate_is_validated_before_io(
+ server: RecordingServer,
+ monkeypatch: pytest.MonkeyPatch,
+ tmp_path: Path,
+ asynchronous: bool,
+ present: bool,
+) -> None:
+ server.expected_requests = 0
+ certificate: Final = tmp_path / "client.pem"
+ if present:
+ certificate.write_text("invalid certificate")
+ monkeypatch.setattr(litellm, "ssl_certificate", str(certificate))
+ with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate.*PEM") as caught:
+ await call(server, asynchronous, num_retries=0)
+ assert str(certificate) not in str(caught.value)
+ assert server.requests == []
diff --git a/tests/test_litellm_rust/responses/__init__.py b/tests/test_litellm_rust/responses/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/test_litellm_rust/responses/conftest.py b/tests/test_litellm_rust/responses/conftest.py
new file mode 100644
index 00000000000..c2ea94473dd
--- /dev/null
+++ b/tests/test_litellm_rust/responses/conftest.py
@@ -0,0 +1,10 @@
+import pytest
+
+from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+
+@pytest.fixture
+def server(recording_server: RecordingServer) -> RecordingServer:
+ recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE)
+ return recording_server
diff --git a/tests/test_litellm_rust/responses/data.py b/tests/test_litellm_rust/responses/data.py
new file mode 100644
index 00000000000..ed489f949a5
--- /dev/null
+++ b/tests/test_litellm_rust/responses/data.py
@@ -0,0 +1,75 @@
+from typing import Final
+
+import litellm
+from litellm.rust_bridge.catalog import Route
+from tests.test_litellm_rust.support.gaps import gap
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import call_natively
+
+RESPONSES_MODEL: Final = "openai/gpt-6.1-sol"
+RESPONSES_TEXT: Final = "Hello from native Responses"
+RESPONSES_RESPONSE: Final = {
+ "id": "resp_native",
+ "object": "response",
+ "created_at": 1,
+ "status": "completed",
+ "model": RESPONSES_MODEL.removeprefix("openai/"),
+ "output": [
+ {
+ "id": "msg_native",
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": RESPONSES_TEXT, "annotations": []}],
+ }
+ ],
+ "usage": {"input_tokens": 5, "output_tokens": 4, "total_tokens": 9},
+}
+RESPONSES_EVENTS: Final = (
+ (
+ "response.created",
+ {
+ "type": "response.created",
+ "sequence_number": 0,
+ "response": {**RESPONSES_RESPONSE, "status": "in_progress", "output": []},
+ },
+ ),
+ (
+ "response.output_text.delta",
+ {
+ "type": "response.output_text.delta",
+ "sequence_number": 1,
+ "item_id": "msg_native",
+ "output_index": 0,
+ "content_index": 0,
+ "delta": RESPONSES_TEXT,
+ },
+ ),
+ ("response.completed", {"type": "response.completed", "sequence_number": 2, "response": RESPONSES_RESPONSE}),
+)
+
+
+def _error(kind: str) -> dict[str, object]:
+ return {"error": {"type": kind, "message": kind, "param": None, "code": None}}
+
+
+REQUEST: Final = {"model": RESPONSES_MODEL, "input": "Hello"}
+BAD_REQUEST: Final = _error("invalid_request_error")
+RATE_LIMITED: Final = _error("rate_limit_exceeded")
+
+
+async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: object) -> object:
+ return await call_natively(
+ Route.RESPONSES,
+ litellm.responses,
+ litellm.aresponses,
+ REQUEST,
+ server,
+ asynchronous=asynchronous,
+ **kwargs,
+ )
+
+
+SYNC_DEPLOYMENT_HOOK: Final = gap("native sync calls skip the async pre-call deployment hook")
+STREAMING: Final = gap("native Responses streaming declines to Python", raises=RuntimeError)
+BODY_HOOKS: Final = gap("native Responses does not run the handler body hooks")
diff --git a/tests/test_litellm_rust/responses/test_body_hooks.py b/tests/test_litellm_rust/responses/test_body_hooks.py
new file mode 100644
index 00000000000..59ba3413da9
--- /dev/null
+++ b/tests/test_litellm_rust/responses/test_body_hooks.py
@@ -0,0 +1,114 @@
+from collections.abc import Iterator
+from datetime import datetime
+from typing import Final
+
+import pytest
+from pydantic import BaseModel
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
+from tests.test_litellm_rust.responses.data import (
+ BODY_HOOKS,
+ REQUEST,
+ RESPONSES_EVENTS,
+ RESPONSES_RESPONSE,
+ STREAMING,
+ call,
+)
+from tests.test_litellm_rust.support.body_hooks import LegacyAgenticLoop, PlannedAgenticLoop, agentic_hooks
+from tests.test_litellm_rust.support.callback_recorder import drain_logging
+from tests.test_litellm_rust.support.isolation import rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import consume, response_id
+from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+STREAM: Final = ResponseSpec(body=None, events=RESPONSES_EVENTS)
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+ yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+@STREAMING
+async def test_body_streaming_deployment_hook_rewrites_chunks_before_delivery(server: RecordingServer) -> None:
+ seen: Final[list[str]] = []
+
+ class Rewrite(CustomLogger):
+ async def async_post_call_streaming_deployment_hook(
+ self, request_data: dict[str, object], response_chunk: BaseModel, call_type: object
+ ) -> BaseModel:
+ seen.append(getattr(response_chunk, "type", ""))
+ if getattr(response_chunk, "type", None) != "response.output_text.delta":
+ return response_chunk
+ return response_chunk.model_copy(update={"delta": "rewritten"})
+
+ server.enqueue(STREAM)
+ with rebound(litellm, "callbacks", [Rewrite()]):
+ stream: Final = await call(server, stream=True)
+ chunks: Final = await consume(stream)
+
+ assert_served_natively(server)
+ assert seen == [event for event, _ in RESPONSES_EVENTS]
+ assert [getattr(chunk, "delta", None) for chunk in chunks] == [None, "rewritten", None]
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_planned_agentic_loop_sends_the_followup_and_replaces_the_response(server: RecordingServer) -> None:
+ loop: Final = PlannedAgenticLoop()
+ server.enqueue(ResponseSpec(body=RESPONSES_RESPONSE))
+ server.default_response = ResponseSpec(body={**RESPONSES_RESPONSE, "id": "followup"})
+
+ with rebound(litellm, "callbacks", [loop]):
+ response: Final = await call(server)
+ await drain_logging()
+
+ assert_served_natively(server, count=2)
+ assert isinstance(server.requests[1].body, dict)
+ assert server.requests[1].body["input"] == [{"role": "user", "content": "followup"}]
+ assert response_id(response) == "agentic-replaced"
+ assert agentic_hooks(loop) == (
+ "async_should_run_agentic_loop",
+ "async_build_agentic_loop_plan",
+ "async_should_run_agentic_loop",
+ "async_post_agentic_loop_response_hook",
+ "async_agentic_loop_cleanup_hook",
+ )
+
+
+@pytest.mark.asyncio
+@BODY_HOOKS
+async def test_body_legacy_agentic_loop_replaces_the_response(server: RecordingServer) -> None:
+ loop: Final = LegacyAgenticLoop()
+
+ with rebound(litellm, "callbacks", [loop]):
+ response: Final = await call(server)
+
+ assert_served_natively(server)
+ assert response_id(response) == "agentic-replaced"
+ assert agentic_hooks(loop) == ("async_should_run_agentic_loop", "async_run_agentic_loop")
+
+
+@pytest.mark.asyncio
+async def test_body_logging_setup_drops_a_logger_fn_from_a_supplied_logging_object(server: RecordingServer) -> None:
+ observed: Final[list[object]] = []
+ logging_obj: Final = Logging(
+ model=REQUEST["model"],
+ messages=[],
+ stream=False,
+ start_time=datetime.now(),
+ call_type="aresponses",
+ litellm_call_id=CALL_ID,
+ function_id=CALL_ID,
+ )
+ logging_obj.update_environment_variables(litellm_params={"logger_fn": observed.append}, optional_params={})
+
+ await call(server, litellm_logging_obj=logging_obj)
+
+ assert_served_natively(server)
+ assert observed == []
diff --git a/tests/test_litellm_rust/responses/test_callbacks.py b/tests/test_litellm_rust/responses/test_callbacks.py
new file mode 100644
index 00000000000..a4abdd43529
--- /dev/null
+++ b/tests/test_litellm_rust/responses/test_callbacks.py
@@ -0,0 +1,408 @@
+import asyncio
+from collections.abc import Iterator
+from contextvars import ContextVar
+from datetime import datetime
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.integrations.argilla import ArgillaLogger
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.integrations.argilla import ArgillaItem
+from litellm.types.utils import StandardLoggingPayload
+from tests.test_litellm_rust.responses.data import (
+ BAD_REQUEST,
+ RATE_LIMITED,
+ RESPONSES_RESPONSE,
+ SYNC_DEPLOYMENT_HOOK,
+ call,
+)
+from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body
+from tests.test_litellm_rust.support.gaps import gap
+from tests.test_litellm_rust.support.isolation import rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import response_id, with_id
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+SYNC_HANDLERS_ON_ASYNC_CALLS: Final = gap("native async calls also run the sync success and failure handlers")
+_OBJECT: Final = TypeAdapter(dict[str, object])
+_HEADERS: Final = TypeAdapter(dict[str, str])
+
+
+def success_signal(asynchronous: bool) -> str:
+ return "async_log_success_event" if asynchronous else "log_success_event"
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+ yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "hook,asynchronous,expected",
+ (
+ pytest.param("log_pre_api_call", True, 1, id="async-log_pre_api_call"),
+ pytest.param("log_post_api_call", True, 1, id="async-log_post_api_call"),
+ pytest.param("async_log_success_event", True, 1, id="async-async_log_success_event"),
+ pytest.param("async_logging_hook", True, 1, id="async-async_logging_hook"),
+ pytest.param("redact_standard_logging_payload_from_model_call_details", True, 1, id="async-redact"),
+ pytest.param("async_pre_call_deployment_hook", True, 1, id="async-async_pre_call_deployment_hook"),
+ pytest.param(
+ "async_post_call_success_deployment_hook", True, 1, id="async-async_post_call_success_deployment_hook"
+ ),
+ pytest.param("log_success_event", True, 0, id="async-log_success_event"),
+ pytest.param("logging_hook", True, 0, id="async-logging_hook"),
+ pytest.param("log_pre_api_call", False, 1, id="sync-log_pre_api_call"),
+ pytest.param("log_post_api_call", False, 1, id="sync-log_post_api_call"),
+ pytest.param("log_success_event", False, 1, id="sync-log_success_event"),
+ pytest.param("logging_hook", False, 1, id="sync-logging_hook"),
+ pytest.param(
+ "async_pre_call_deployment_hook",
+ False,
+ 1,
+ id="sync-async_pre_call_deployment_hook",
+ marks=SYNC_DEPLOYMENT_HOOK,
+ ),
+ pytest.param(
+ "async_post_call_success_deployment_hook", False, 0, id="sync-async_post_call_success_deployment_hook"
+ ),
+ ),
+)
+async def test_success_hook_fires_the_expected_number_of_times(
+ server: RecordingServer, hook: str, asynchronous: bool, expected: int
+) -> None:
+ recorder: Final = HookRecorder()
+
+ with rebound(litellm, "callbacks", [recorder]):
+ await call(server, asynchronous=asynchronous)
+ await recorder.wait_for_async(success_signal(asynchronous))
+
+ assert_served_natively(server)
+ assert recorder.names.count(hook) == expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "hook,asynchronous,expected",
+ (
+ pytest.param("async_log_failure_event", True, 1, id="async-async_log_failure_event"),
+ pytest.param(
+ "async_post_call_failure_deployment_hook", True, 1, id="async-async_post_call_failure_deployment_hook"
+ ),
+ pytest.param("log_model_group_rate_limit_error", True, 1, id="async-log_model_group_rate_limit_error"),
+ pytest.param("log_failure_event", True, 0, id="async-log_failure_event", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+ pytest.param("log_failure_event", False, 1, id="sync-log_failure_event"),
+ ),
+)
+async def test_failure_hook_fires_the_expected_number_of_times(
+ server: RecordingServer, hook: str, asynchronous: bool, expected: int
+) -> None:
+ server.default_response = ResponseSpec(body=RATE_LIMITED, status=429)
+ recorder: Final = HookRecorder()
+
+ with rebound(litellm, "callbacks", [recorder]):
+ with pytest.raises(litellm.RateLimitError):
+ await call(server, asynchronous=asynchronous, metadata={"model_group_size": 1})
+ await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event")
+
+ assert_served_natively(server)
+ assert recorder.names.count(hook) == expected
+
+
+@pytest.mark.asyncio
+async def test_failure_events_carry_the_error_raised_to_the_caller(server: RecordingServer) -> None:
+ server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
+ recorder: Final = RecordingLogger()
+
+ with rebound(litellm, "callbacks", [recorder]):
+ with pytest.raises(litellm.BadRequestError) as raised:
+ await call(server)
+ events: Final = await recorder.wait_for_async("async_log_failure_event")
+
+ assert_served_natively(server)
+ assert isinstance(events[0].kwargs, dict)
+ assert events[0].kwargs["exception"] is raised.value
+
+
+@pytest.mark.asyncio
+async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None:
+ recorder: Final = RecordingLogger()
+
+ await call(server, callbacks=[recorder])
+
+ assert_served_natively(server)
+ pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs
+ assert isinstance(pre_call, dict)
+ assert request_body(pre_call) == server.requests[0].body
+
+
+@pytest.mark.asyncio
+async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None:
+ class Edit(CustomLogger):
+ def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
+ request_body(kwargs)["temperature"] = 0.25
+
+ await call(server, callbacks=[Edit()])
+
+ assert_served_natively(server)
+ assert isinstance(server.requests[0].body, dict)
+ assert server.requests[0].body["temperature"] == 0.25
+
+
+@pytest.mark.asyncio
+async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None:
+ class Rewrite(HookRecorder):
+ async def async_logging_hook(
+ self, kwargs: dict[str, object], result: object, call_type: str
+ ) -> tuple[dict[str, object], object]:
+ return {**kwargs, "logging_hook_ran": True}, result
+
+ def redact_standard_logging_payload_from_model_call_details(
+ self, model_call_details: dict[str, object]
+ ) -> dict[str, object]:
+ return {**model_call_details, "redacted": True}
+
+ recorder: Final = Rewrite()
+
+ with rebound(litellm, "callbacks", [recorder]):
+ await call(server)
+ success: Final = await recorder.wait_for_async("async_log_success_event")
+
+ assert isinstance(success[0].kwargs, dict)
+ assert success[0].kwargs["logging_hook_ran"] is True
+ assert success[0].kwargs["redacted"] is True
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_hooks_compose_the_wire_request_in_registration_order(
+ server: RecordingServer,
+) -> None:
+ class Rewrite(CustomLogger):
+ def __init__(self, marker: str) -> None:
+ super().__init__()
+ self.marker = marker
+
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: object
+ ) -> dict[str, object]:
+ headers: Final = _HEADERS.validate_python(kwargs["extra_headers"])
+ return {
+ **kwargs,
+ "extra_headers": {**headers, "x-deployment-hook": headers["x-deployment-hook"] + self.marker},
+ }
+
+ with rebound(litellm, "callbacks", [Rewrite("first"), Rewrite("-second")]):
+ await call(server, extra_headers={"x-deployment-hook": "caller-", "x-kept": "unchanged"})
+
+ assert_served_natively(server)
+ assert server.requests[0].headers["x-deployment-hook"] == "caller-first-second"
+ assert server.requests[0].headers["x-kept"] == "unchanged"
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(server: RecordingServer) -> None:
+ rejection: Final = litellm.BadRequestError(message="deployment rejected", model="test-model", llm_provider="openai")
+ recorder: Final = RecordingLogger()
+
+ class Reject(CustomLogger):
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: object
+ ) -> dict[str, object]:
+ raise rejection
+
+ with rebound(litellm, "callbacks", [Reject(), recorder]):
+ with pytest.raises(litellm.BadRequestError, match="deployment rejected") as raised:
+ await call(server)
+ failures: Final = await recorder.wait_for_async("async_log_failure_event")
+
+ assert raised.value is rejection
+ assert server.requests == []
+ assert len(failures) == 1
+ assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value
+ assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names)
+
+
+@pytest.mark.asyncio
+async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None:
+ class DropKeywords(CustomLogger):
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: object
+ ) -> dict[str, object]:
+ return {name: value for name, value in kwargs.items() if name not in ("instructions", "extra_headers")}
+
+ with rebound(litellm, "callbacks", [DropKeywords()]):
+ await call(
+ server,
+ instructions="remove this instruction",
+ extra_headers={"x-deleted": "remove this header"},
+ )
+
+ assert_served_natively(server)
+ sent: Final = server.requests[0]
+ assert isinstance(sent.body, dict)
+ assert "instructions" not in sent.body
+ assert "x-deleted" not in sent.headers
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_deployment_hook_replaces_the_response(server: RecordingServer) -> None:
+ recorder: Final = RecordingLogger()
+
+ class Replace(CustomLogger):
+ async def async_post_call_success_deployment_hook(
+ self, request_data: dict[str, object], response: object, call_type: object
+ ) -> object:
+ return with_id(response, "replaced")
+
+ with rebound(litellm, "callbacks", [Replace(), recorder]):
+ response: Final = await call(server)
+ success: Final = await recorder.wait_for_async("async_log_success_event")
+
+ assert_served_natively(server)
+ assert response_id(response) == "replaced"
+ assert len(success) == 1
+ assert response_id(success[0].response) == response_id(response)
+
+
+@pytest.mark.asyncio
+async def test_concurrent_calls_keep_callback_context_metadata_and_responses_isolated(server: RecordingServer) -> None:
+ context: Final = ContextVar("responses-callback-context", default="outside")
+ first_entered: Final = asyncio.Event()
+ second_entered: Final = asyncio.Event()
+ server.enqueue(ResponseSpec(body={**RESPONSES_RESPONSE, "id": "resp_first"}))
+ server.enqueue(ResponseSpec(body={**RESPONSES_RESPONSE, "id": "resp_second"}))
+
+ class Observe(RecordingLogger):
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: object
+ ) -> dict[str, object]:
+ arrived, other = (
+ (first_entered, second_entered)
+ if kwargs["litellm_call_id"] == "first"
+ else (second_entered, first_entered)
+ )
+ arrived.set()
+ await asyncio.wait_for(other.wait(), timeout=5)
+ return kwargs
+
+ def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
+ self._record("log_pre_api_call", {**kwargs, "caller_context": context.get()})
+
+ async def async_log_success_event(
+ self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+ ) -> None:
+ self._record("async_log_success_event", {**kwargs, "caller_context": context.get()}, response_obj)
+
+ recorder: Final = Observe()
+
+ async def invoke(marker: str) -> object:
+ token: Final = context.set(marker)
+ try:
+ return await call(server, litellm_call_id=marker, metadata={"marker": marker})
+ finally:
+ context.reset(token)
+
+ with rebound(litellm, "callbacks", [recorder]):
+ responses: Final = await asyncio.gather(invoke("first"), invoke("second"))
+ success: Final = await recorder.wait_for_async("async_log_success_event", count=2)
+
+ assert_served_natively(server, count=2)
+ expected: Final = {"first": responses[0], "second": responses[1]}
+ pre_call: Final = recorder.wait_for("log_pre_api_call", count=2)
+ assert len(pre_call) == len(success) == 2
+ assert {_OBJECT.validate_python(event.kwargs)["litellm_call_id"] for event in success} == expected.keys()
+ for event in (*pre_call, *success):
+ details: Final = _OBJECT.validate_python(event.kwargs)
+ call_id: Final = str(details["litellm_call_id"])
+ params: Final = _OBJECT.validate_python(details["litellm_params"])
+ assert details["caller_context"] == call_id
+ assert _OBJECT.validate_python(params["metadata"])["marker"] == call_id
+ for event in success:
+ logged_call_id: Final = str(_OBJECT.validate_python(event.kwargs)["litellm_call_id"])
+ assert response_id(event.response) == response_id(expected[logged_call_id])
+ assert context.get() == "outside"
+
+
+@pytest.mark.asyncio
+async def test_a_failing_logger_does_not_fail_the_call(server: RecordingServer) -> None:
+ class Broken(CustomLogger):
+ def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
+ raise RuntimeError("logger exploded")
+
+ response: Final = await call(server, callbacks=[Broken()])
+
+ assert_served_natively(server)
+ assert response_id(response) == RESPONSES_RESPONSE["id"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "event,asynchronous,dynamic,expected",
+ (
+ pytest.param("input", False, False, 1, id="global-input"),
+ pytest.param("success", True, False, 1, id="global-async-success"),
+ pytest.param("success", True, True, 1, id="request-async-success"),
+ pytest.param("success", False, False, 0, id="global-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+ pytest.param("success", False, True, 0, id="request-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
+ pytest.param("failure", True, False, 1, id="global-async-failure"),
+ pytest.param("failure", False, False, 1, id="global-sync-failure"),
+ pytest.param("failure", False, True, 1, id="request-sync-failure"),
+ ),
+)
+async def test_function_callback_fires_the_expected_number_of_times(
+ server: RecordingServer, event: str, asynchronous: bool, dynamic: bool, expected: int
+) -> None:
+ recorder: Final = RecordingLogger()
+
+ def callback(kwargs: dict[str, object], *_: object) -> None:
+ recorder._record(event, kwargs)
+
+ async def async_callback(kwargs: dict[str, object], *_: object) -> None:
+ recorder._record(event, kwargs)
+
+ registered: Final = async_callback if asynchronous else callback
+ global_list: Final = f"_async_{event}_callback" if asynchronous else f"{event}_callback"
+ request_callbacks: Final = {f"{event}_callback": [registered]} if dynamic else {}
+ if event == "failure":
+ server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
+
+ with rebound(litellm, "callbacks", [recorder]), rebound(litellm, global_list, [] if dynamic else [registered]):
+ try:
+ await call(server, **request_callbacks)
+ except litellm.BadRequestError:
+ assert event == "failure"
+ await recorder.wait_for_async(f"async_log_{'failure' if event == 'failure' else 'success'}_event")
+
+ assert_served_natively(server)
+ assert recorder.names.count(event) == expected
+
+
+@pytest.mark.asyncio
+async def test_dataset_hook_returning_none_drops_the_record_from_the_success_sink(server: RecordingServer) -> None:
+ datasets: Final = RecordingLogger()
+
+ class Drop(CustomLogger):
+ async def async_dataset_hook(
+ self, logged_data: ArgillaItem, standard_logging_payload: StandardLoggingPayload | None
+ ) -> None:
+ datasets._record("async_dataset_hook", response=logged_data)
+
+ server.enqueue(ResponseSpec(body={"items": [{"id": "callback-dataset"}]}))
+ with rebound(litellm, "argilla_transformation_object", {"prompt": "messages", "answer": "response"}):
+ sink: Final = ArgillaLogger(
+ argilla_api_key="test-key",
+ argilla_base_url=server.base_url,
+ argilla_dataset_name="callback-dataset",
+ flush_interval=3600,
+ )
+ with rebound(litellm, "callbacks", [sink, Drop()]):
+ await call(server)
+ await datasets.wait_for_async("async_dataset_hook")
+ await sink.flush_queue()
+
+ assert [request.method for request in server.requests] == ["GET", "POST"]
diff --git a/tests/test_litellm_rust/responses/test_failures.py b/tests/test_litellm_rust/responses/test_failures.py
new file mode 100644
index 00000000000..6dad1bc7278
--- /dev/null
+++ b/tests/test_litellm_rust/responses/test_failures.py
@@ -0,0 +1,83 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm import RateLimitError
+from tests.test_litellm_rust.responses.data import RATE_LIMITED, call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+def failure_signal(asynchronous: bool) -> str:
+ return "async_log_failure_event" if asynchronous else "log_failure_event"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool
+) -> None:
+ monkeypatch.setattr(litellm, "max_budget", 1.0)
+ monkeypatch.setattr(litellm, "_current_cost", 2.0)
+ monkeypatch.setattr(litellm, "ssl_verify", object())
+ recorder: Final = RecordingLogger()
+ server.expected_requests = 0
+
+ with pytest.raises(litellm.BudgetExceededError) as caught:
+ await call(server, asynchronous=asynchronous, callbacks=[recorder])
+
+ failure: Final = await recorder.wait_for_async(failure_signal(asynchronous))
+ assert len(failure) == 1
+ assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+ assert not server.requests
+ assert not any("success" in name for name in recorder.names)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_provider_failure_is_terminal_and_shared_with_callbacks(
+ server: RecordingServer, asynchronous: bool
+) -> None:
+ recorder: Final = RecordingLogger()
+ server.enqueue(ResponseSpec(body=RATE_LIMITED, status=429))
+
+ with pytest.raises(RateLimitError) as caught:
+ await call(server, asynchronous=asynchronous, callbacks=[recorder])
+
+ assert caught.value.status_code == 429
+ assert len(server.requests) == 1
+ failure: Final = await recorder.wait_for_async(failure_signal(asynchronous))
+ assert len(failure) == 1
+ assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+ assert not any("success" in name for name in recorder.names)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "options",
+ (
+ {"stream": True},
+ {"extra_body": {"provider_option": True}},
+ {"mock_response": "mock"},
+ {"num_retries": 1},
+ {"use_chat_completions_api": True},
+ {"model_list": []},
+ ),
+)
+async def test_unsupported_requests_fail_before_callbacks_and_provider_io(
+ server: RecordingServer, options: dict[str, object]
+) -> None:
+ recorder: Final = RecordingLogger()
+ server.expected_requests = 0
+
+ with pytest.raises(RuntimeError, match="declined the request"):
+ await call(server, callbacks=[recorder], **options)
+
+ assert not server.requests
+ assert not recorder.events
diff --git a/tests/test_litellm_rust/responses/test_lifecycle.py b/tests/test_litellm_rust/responses/test_lifecycle.py
new file mode 100644
index 00000000000..a8a4ce18f09
--- /dev/null
+++ b/tests/test_litellm_rust/responses/test_lifecycle.py
@@ -0,0 +1,46 @@
+import asyncio
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.rust_bridge import _native
+from litellm.rust_bridge.public_call import NativeCall
+from tests.test_litellm_rust.responses.data import REQUEST
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.mark.asyncio
+async def test_unstarted_call_has_no_provider_or_callback_effects(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorder: Final = RecordingLogger()
+ server.expected_requests = 0
+ monkeypatch.setattr(litellm, "ssl_verify", object())
+ monkeypatch.setattr(litellm, "max_budget", 1.0)
+ monkeypatch.setattr(litellm, "_current_cost", 2.0)
+ kwargs: Final = {
+ **REQUEST,
+ "api_key": "test-key",
+ "api_base": server.base_url,
+ "callbacks": [recorder],
+ }
+
+ base: Final = {
+ **REQUEST,
+ "stream": None,
+ "api_key": "test-key",
+ "api_base": server.base_url,
+ "custom_llm_provider": "openai",
+ "extra_headers": None,
+ }
+
+ pending: Final = _native.aresponses(NativeCall(args=(), kwargs=kwargs, base=base))
+ assert asyncio.iscoroutine(pending)
+ pending.close()
+
+ assert not server.requests
+ assert not recorder.events
diff --git a/tests/test_litellm_rust/responses/test_requests.py b/tests/test_litellm_rust/responses/test_requests.py
new file mode 100644
index 00000000000..db04ccb5c65
--- /dev/null
+++ b/tests/test_litellm_rust/responses/test_requests.py
@@ -0,0 +1,56 @@
+import asyncio
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.responses.utils import ResponsesAPIRequestUtils
+from litellm.rust_bridge.catalog import Route
+from litellm.types.llms.openai import ResponsesAPIResponse
+from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_TEXT, call
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively, native_route
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+async def test_positional_parameters_reach_the_wire(server: RecordingServer, asynchronous: bool) -> None:
+ arguments: Final = ("hello", RESPONSES_MODEL, None, "Be brief", 16)
+ supplied: Final = {"api_key": "test-key", "api_base": server.base_url}
+
+ with native_route(Route.RESPONSES):
+ result: Final = (
+ await litellm.aresponses(*arguments, **supplied)
+ if asynchronous
+ else await asyncio.to_thread(litellm.responses, *arguments, **supplied)
+ )
+
+ assert isinstance(result, ResponsesAPIResponse)
+ assert result.output_text == RESPONSES_TEXT
+ assert_served_natively(server)
+ body: Final = _OBJECT.validate_python(server.requests[0].body)
+ assert body["instructions"] == arguments[3]
+ assert body["max_output_tokens"] == arguments[4]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+@pytest.mark.parametrize("encoded", (False, True))
+async def test_continuation_ids_are_decoded_before_they_reach_the_wire(
+ server: RecordingServer, asynchronous: bool, encoded: bool
+) -> None:
+ original: Final = "resp_upstream"
+ previous: Final = (
+ ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "deployment", original)
+ if encoded
+ else original
+ )
+
+ await call(server, asynchronous=asynchronous, previous_response_id=previous)
+
+ assert _OBJECT.validate_python(server.requests[0].body)["previous_response_id"] == original
diff --git a/tests/test_litellm_rust/responses/test_settings.py b/tests/test_litellm_rust/responses/test_settings.py
new file mode 100644
index 00000000000..883a84b2a45
--- /dev/null
+++ b/tests/test_litellm_rust/responses/test_settings.py
@@ -0,0 +1,83 @@
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.models.credentials import CredentialItem
+from litellm.types.llms.openai import ResponsesAPIResponse
+from litellm.types.utils import CallTypes
+from tests.test_litellm_rust.responses.data import call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+_OBJECT: Final = TypeAdapter(dict[str, object])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("from_credentials", (False, True))
+async def test_resource_setup_uses_deployment_hook_arguments(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, from_credentials: bool
+) -> None:
+ invalid_settings: Final = {"ssl_verify": object()}
+ credential: Final = CredentialItem(
+ credential_name="resource-settings", credential_info={}, credential_values=invalid_settings
+ )
+ monkeypatch.setattr(litellm, "credential_list", [credential])
+
+ class Prepare(CustomLogger):
+ async def async_pre_call_deployment_hook(
+ self, kwargs: dict[str, object], call_type: CallTypes | None
+ ) -> dict[str, object]:
+ return {
+ **kwargs,
+ **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings),
+ }
+
+ litellm.callbacks.append(Prepare())
+ recorder: Final = RecordingLogger()
+ server.expected_requests = 0
+
+ with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught:
+ await call(server, callbacks=[recorder])
+
+ failure: Final = await recorder.wait_for_async("async_log_failure_event")
+ assert len(failure) == 1
+ assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
+ assert not server.requests
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", (False, True))
+@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty"))
+async def test_connection_settings_reach_the_provider(
+ server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, source: str
+) -> None:
+ key: Final = "selected-key"
+ explicit: Final = source in ("explicit", "base_url")
+ monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None))
+ monkeypatch.setattr(litellm, "openai_key", key if source == "provider" else ("unused" if explicit else None))
+ monkeypatch.setattr(
+ litellm,
+ "api_base",
+ None if source == "environment" else ("http://127.0.0.1:1" if explicit else server.base_url),
+ )
+ monkeypatch.setenv("OPENAI_API_KEY", key if source == "environment" else "unused")
+ for name in ("OPENAI_BASE_URL", "OPENAI_API_BASE"):
+ monkeypatch.setenv(name, server.base_url if source == "environment" else "http://127.0.0.1:1")
+
+ result: Final = await call(
+ server,
+ asynchronous=asynchronous,
+ api_key=key if explicit else ("" if source == "empty" else None),
+ api_base=server.base_url if source == "explicit" else ("" if source == "empty" else None),
+ **({"base_url": server.base_url} if source == "base_url" else {}),
+ )
+
+ assert isinstance(result, ResponsesAPIResponse)
+ assert_served_natively(server)
+ assert server.requests[0].headers["authorization"] == f"Bearer {key}"
diff --git a/tests/test_litellm_rust/responses/test_streaming.py b/tests/test_litellm_rust/responses/test_streaming.py
new file mode 100644
index 00000000000..fe284bfa3fc
--- /dev/null
+++ b/tests/test_litellm_rust/responses/test_streaming.py
@@ -0,0 +1,58 @@
+from collections.abc import AsyncIterator, Iterator
+from typing import Final
+
+import pytest
+
+from tests.test_litellm_rust.responses.data import RESPONSES_EVENTS, RESPONSES_RESPONSE, STREAMING, call
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
+from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
+from tests.test_litellm_rust.support.results import consume, response_id
+from tests.test_litellm_rust.support.routing import assert_served_natively
+
+pytestmark = pytest.mark.requires_rust_extension
+
+STREAM: Final = ResponseSpec(body=None, events=RESPONSES_EVENTS)
+
+
+@pytest.fixture
+def server(server: RecordingServer) -> Iterator[RecordingServer]:
+ yield from relaxed(server)
+
+
+@pytest.mark.asyncio
+@STREAMING
+async def test_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk(
+ server: RecordingServer,
+) -> None:
+ server.enqueue(STREAM)
+ recorder: Final = RecordingLogger()
+
+ stream: Final = await call(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_served_natively(server)
+ assert [getattr(chunk, "type", None) for chunk in (first, *rest)] == [event for event, _ in RESPONSES_EVENTS]
+ success: Final = await recorder.wait_for_async("async_log_success_event")
+ assert len(success) == 1
+ assert isinstance(success[0].kwargs, dict)
+ assert success[0].kwargs["completion_start_time"] is not None
+ assert response_id(success[0].response) == RESPONSES_RESPONSE["id"]
+ assert "log_failure_event" not in recorder.names
+
+
+@pytest.mark.asyncio
+@STREAMING
+async def test_sync_stream_relays_provider_events_and_logs_success_once(server: RecordingServer) -> None:
+ server.enqueue(STREAM)
+ recorder: Final = RecordingLogger()
+
+ stream: Final = await call(server, asynchronous=False, stream=True, callbacks=[recorder])
+ chunks: Final = await consume(stream)
+
+ assert_served_natively(server)
+ assert [getattr(chunk, "type", None) for chunk in chunks] == [event for event, _ in RESPONSES_EVENTS]
+ assert len(await recorder.wait_for_async("log_success_event")) == 1
diff --git a/tests/test_litellm_rust/support/body_hooks.py b/tests/test_litellm_rust/support/body_hooks.py
new file mode 100644
index 00000000000..1764b84da67
--- /dev/null
+++ b/tests/test_litellm_rust/support/body_hooks.py
@@ -0,0 +1,39 @@
+from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch
+from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
+from tests.test_litellm_rust.support.results import response_id, with_id
+
+
+def agentic_hooks(recorder: RecordingLogger) -> tuple[str, ...]:
+ return tuple(name for name in recorder.names if "agentic" in name)
+
+
+class PlannedAgenticLoop(RecordingLogger):
+ async def async_should_run_agentic_loop(self, response: object, **kwargs: object) -> tuple[bool, dict[str, object]]:
+ self._record("async_should_run_agentic_loop", response=response)
+ return response_id(response) != "followup", {}
+
+ async def async_build_agentic_loop_plan(self, response: object, **kwargs: object) -> AgenticLoopPlan:
+ self._record("async_build_agentic_loop_plan", response=response)
+ return AgenticLoopPlan(
+ run_agentic_loop=True,
+ request_patch=AgenticLoopRequestPatch.model_validate(
+ {"messages": [{"role": "user", "content": "followup"}]}
+ ),
+ )
+
+ async def async_post_agentic_loop_response_hook(self, response: object, **kwargs: object) -> object:
+ self._record("async_post_agentic_loop_response_hook", response=response)
+ return with_id(response, "agentic-replaced")
+
+ async def async_agentic_loop_cleanup_hook(self, **kwargs: object) -> None:
+ self._record("async_agentic_loop_cleanup_hook")
+
+
+class LegacyAgenticLoop(RecordingLogger):
+ async def async_should_run_agentic_loop(self, response: object, **kwargs: object) -> tuple[bool, dict[str, object]]:
+ self._record("async_should_run_agentic_loop", response=response)
+ return True, {}
+
+ async def async_run_agentic_loop(self, response: object, **kwargs: object) -> object:
+ self._record("async_run_agentic_loop", response=response)
+ return with_id(response, "agentic-replaced")
diff --git a/tests/test_litellm_rust/support/callback_recorder.py b/tests/test_litellm_rust/support/callback_recorder.py
index d3749ccc095..c5e9fcd51a8 100644
--- a/tests/test_litellm_rust/support/callback_recorder.py
+++ b/tests/test_litellm_rust/support/callback_recorder.py
@@ -3,6 +3,7 @@ import copy
import threading
import time
from dataclasses import dataclass
+from datetime import datetime
from typing import Final
from litellm.integrations.custom_logger import CustomLogger
@@ -99,3 +100,63 @@ class RecordingLogger(CustomLogger):
def logging_hook(self, kwargs, result, call_type):
self._record("logging_hook", kwargs, result)
return kwargs, result
+
+
+class HookRecorder(RecordingLogger):
+ def log_post_api_call(
+ self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+ ) -> None:
+ self._record("log_post_api_call", kwargs, response_obj)
+
+ async def async_logging_hook(
+ self, kwargs: dict[str, object], result: object, call_type: str
+ ) -> tuple[dict[str, object], object]:
+ self._record("async_logging_hook", kwargs, result)
+ return kwargs, result
+
+ def redact_standard_logging_payload_from_model_call_details(
+ self, model_call_details: dict[str, object]
+ ) -> dict[str, object]:
+ self._record("redact_standard_logging_payload_from_model_call_details", model_call_details)
+ return model_call_details
+
+ async def async_pre_call_deployment_hook(self, kwargs: dict[str, object], call_type: object) -> dict[str, object]:
+ self._record("async_pre_call_deployment_hook", kwargs)
+ return kwargs
+
+ async def async_post_call_success_deployment_hook(
+ self, request_data: dict[str, object], response: object, call_type: object
+ ) -> object:
+ self._record("async_post_call_success_deployment_hook", request_data, response)
+ return response
+
+ async def async_post_call_failure_deployment_hook(
+ self,
+ request_data: dict[str, object],
+ original_exception: Exception,
+ call_type: object,
+ *,
+ fallback_depth: int | None = None,
+ ) -> None:
+ self._record("async_post_call_failure_deployment_hook", request_data, original_exception)
+
+ async def log_model_group_rate_limit_error(
+ self, exception: Exception, original_model_group: str | None, kwargs: dict[str, object]
+ ) -> None:
+ self._record("log_model_group_rate_limit_error", kwargs, exception)
+
+
+def request_body(kwargs: dict[str, object]) -> dict[str, object]:
+ additional_args = kwargs["additional_args"]
+ assert isinstance(additional_args, dict)
+ body = additional_args["complete_input_dict"]
+ assert isinstance(body, dict)
+ return body
+
+
+def request_headers(kwargs: dict[str, object]) -> dict[str, object]:
+ additional_args = kwargs["additional_args"]
+ assert isinstance(additional_args, dict)
+ headers = additional_args["headers"]
+ assert isinstance(headers, dict)
+ return headers
diff --git a/tests/test_litellm_rust/support/gaps.py b/tests/test_litellm_rust/support/gaps.py
new file mode 100644
index 00000000000..bd2cefec467
--- /dev/null
+++ b/tests/test_litellm_rust/support/gaps.py
@@ -0,0 +1,5 @@
+import pytest
+
+
+def gap(reason: str, raises: type[Exception] = AssertionError) -> pytest.MarkDecorator:
+ return pytest.mark.xfail(reason=reason, raises=raises, strict=True)
diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py
index 74ca2cda1c5..d7b82767b97 100644
--- a/tests/test_litellm_rust/support/recording_server.py
+++ b/tests/test_litellm_rust/support/recording_server.py
@@ -120,3 +120,9 @@ def recording_service() -> Iterator[RecordingServer]:
if recording_server.expected_requests is not None:
assert len(recording_server.requests) == recording_server.expected_requests
assert recording_server.responses == []
+
+
+def relaxed(server: RecordingServer) -> Iterator[RecordingServer]:
+ server.expected_requests = None
+ yield server
+ server.responses.clear()
diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py
deleted file mode 100644
index c9cf81b83ca..00000000000
--- a/tests/test_litellm_rust/support/requests.py
+++ /dev/null
@@ -1,97 +0,0 @@
-from typing import Final
-
-import litellm
-from litellm.llms.base_llm.ocr.transformation import OCRResponse
-from tests.test_litellm_rust.support.recording_server import RecordingServer
-
-OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}
-OCR_MODEL: Final = "mistral/mistral-ocr-latest"
-OCR_RESPONSE: Final = {
- "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
- "model": "mistral-ocr-latest",
- "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 {
- "model": OCR_MODEL,
- "document": dict(OCR_DOCUMENT),
- "api_key": "test-key",
- "api_base": server.base_url,
- **kwargs,
- }
-
-
-def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
- response: Final = litellm.ocr(**ocr_arguments(server, **kwargs))
- if not isinstance(response, OCRResponse):
- raise TypeError(f"Expected OCRResponse, got {type(response).__name__}")
- return response
-
-
-async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
- return await litellm.aocr(**ocr_arguments(server, **kwargs))
-
-
-def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
- return call_ocr(server, **kwargs)
-
-
-async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
- return await call_aocr(server, **kwargs)
-
-
-async def call_native(server: RecordingServer, asynchronous: bool, **kwargs: object) -> OCRResponse:
- return await call_native_aocr(server, **kwargs) if asynchronous else call_native_ocr(server, **kwargs)
-
-
-def request_body(kwargs: dict[str, object]) -> dict[str, object]:
- additional_args = kwargs["additional_args"]
- assert isinstance(additional_args, dict)
- body = additional_args["complete_input_dict"]
- assert isinstance(body, dict)
- return body
-
-
-def request_headers(kwargs: dict[str, object]) -> dict[str, object]:
- additional_args = kwargs["additional_args"]
- assert isinstance(additional_args, dict)
- headers = additional_args["headers"]
- assert isinstance(headers, dict)
- return headers
diff --git a/tests/test_litellm_rust/support/results.py b/tests/test_litellm_rust/support/results.py
new file mode 100644
index 00000000000..814cebe3efc
--- /dev/null
+++ b/tests/test_litellm_rust/support/results.py
@@ -0,0 +1,26 @@
+from collections.abc import AsyncIterator, Iterator
+
+from pydantic import BaseModel
+
+from litellm.responses.utils import ResponsesAPIRequestUtils
+from litellm.types.llms.openai import ResponsesAPIResponse
+
+
+def response_id(response: object) -> object:
+ if isinstance(response, ResponsesAPIResponse):
+ return ResponsesAPIRequestUtils.decode_responses_api_response_id(response.id).get("response_id") or response.id
+ return response.get("id") if isinstance(response, dict) else None
+
+
+def with_id(response: object, new_id: str) -> object:
+ if isinstance(response, BaseModel):
+ return response.model_copy(update={"id": new_id})
+ assert isinstance(response, dict)
+ return {**response, "id": new_id}
+
+
+async def consume(stream: object) -> tuple[object, ...]:
+ if isinstance(stream, AsyncIterator):
+ return tuple([chunk async for chunk in stream])
+ assert isinstance(stream, Iterator)
+ return tuple(stream)
diff --git a/tests/test_litellm_rust/support/routing.py b/tests/test_litellm_rust/support/routing.py
new file mode 100644
index 00000000000..56bb1276012
--- /dev/null
+++ b/tests/test_litellm_rust/support/routing.py
@@ -0,0 +1,48 @@
+import asyncio
+import copy
+from collections.abc import Awaitable, Callable, Generator, Mapping
+from contextlib import contextmanager
+from typing import Final
+
+from litellm.rust_bridge import catalog
+from litellm.rust_bridge.catalog import Route, RouteRule
+from litellm.rust_bridge.configuration import Rollout
+from tests.test_litellm_rust.support.isolation import rebound
+from tests.test_litellm_rust.support.recording_server import RecordingServer
+from tests.test_litellm_rust.support.served import assert_served_natively
+
+__all__ = ["CALL_ID", "assert_served_natively", "call_natively", "native_route"]
+
+CALL_ID: Final = "callback-contract"
+
+
+@contextmanager
+def native_route(*routes: Route) -> Generator[None]:
+ rules: tuple[RouteRule, ...] = tuple(RouteRule(route, Rollout.RUST_REQUIRED) for route in routes)
+ with rebound(catalog, "RULES", (*rules, *catalog.RULES)):
+ yield
+
+
+async def call_natively(
+ route: Route,
+ create: Callable[..., object],
+ acreate: Callable[..., Awaitable[object]],
+ request: Mapping[str, object],
+ server: RecordingServer,
+ *,
+ asynchronous: bool,
+ **kwargs: object,
+) -> object:
+ arguments: Final = {
+ **copy.deepcopy(dict(request)),
+ "api_key": "test-key",
+ "api_base": server.base_url,
+ "litellm_call_id": CALL_ID,
+ **kwargs,
+ }
+ with native_route(route):
+ if asynchronous:
+ return await acreate(**arguments)
+ return await asyncio.to_thread(create, **arguments)
+
+
diff --git a/tests/test_litellm_rust/support/served.py b/tests/test_litellm_rust/support/served.py
index d3c568937f3..edf492ea7e6 100644
--- a/tests/test_litellm_rust/support/served.py
+++ b/tests/test_litellm_rust/support/served.py
@@ -4,8 +4,13 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
from tests.test_litellm_rust.support.recording_server import RecordingServer
-def assert_served_natively(server: RecordingServer, response: object) -> None:
- assert len(server.requests) == 1
- additional_headers: Final = get_hidden_params_dict(response).get("additional_headers")
- assert isinstance(additional_headers, dict)
- assert additional_headers.get("x-litellm-rust") == "true"
+def assert_served_natively(server: RecordingServer, *responses: object, count: int = 1) -> None:
+ assert len(server.requests) == (len(responses) or count)
+ if not responses:
+ assert not any(
+ request.headers.get("user-agent", "").startswith("python-httpx") for request in server.requests
+ )
+ for response in responses:
+ additional_headers: Final = get_hidden_params_dict(response).get("additional_headers")
+ assert isinstance(additional_headers, dict)
+ assert additional_headers.get("x-litellm-rust") == "true"
diff --git a/tests/test_litellm_rust/test_inference.py b/tests/test_litellm_rust/test_inference.py
deleted file mode 100644
index 83a01e8ad4d..00000000000
--- a/tests/test_litellm_rust/test_inference.py
+++ /dev/null
@@ -1,403 +0,0 @@
-import asyncio
-from collections.abc import Awaitable, Coroutine, Mapping
-from typing import Final, Literal, TypeAlias
-
-import pytest
-from pydantic import JsonValue, TypeAdapter
-
-import litellm
-from litellm import RateLimitError
-from litellm.chat_completions import dispatch as chat_dispatch
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.models.credentials import CredentialItem
-from litellm.responses.utils import ResponsesAPIRequestUtils
-from litellm.rust_bridge import _native
-from litellm.rust_bridge.public_call import NativeCall
-from litellm.types.llms.openai import ResponsesAPIResponse
-from litellm.types.utils import CallTypes, ModelResponse
-from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
-from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE, request_body
-
-pytestmark = pytest.mark.requires_rust_extension
-Route: TypeAlias = Literal["chat", "responses"]
-_OBJECT: Final = TypeAdapter(dict[str, object])
-NativeResult: TypeAlias = (
- ModelResponse
- | ResponsesAPIResponse
- | Coroutine[object, object, ModelResponse]
- | Coroutine[object, object, ResponsesAPIResponse]
-)
-RESPONSES_MODEL: Final = "openai/gpt-6-sol"
-RESPONSES_RESPONSE: Final[dict[str, JsonValue]] = {
- "id": "resp_native",
- "object": "response",
- "created_at": 1,
- "model": RESPONSES_MODEL.removeprefix("openai/"),
- "status": "completed",
- "output": [
- {
- "type": "message",
- "id": "msg_native",
- "role": "assistant",
- "status": "completed",
- "content": [{"type": "output_text", "text": "native response", "annotations": []}],
- }
- ],
- "usage": {"input_tokens": 5, "output_tokens": 4, "total_tokens": 9},
-}
-
-
-@pytest.fixture(params=("chat", "responses"))
-def route(request: pytest.FixtureRequest) -> Route:
- return TypeAdapter(Route).validate_python(request.param)
-
-
-def native_call(
- route: Route, asynchronous: bool, server: RecordingServer, options: Mapping[str, object]
-) -> NativeResult:
- server.default_response = ResponseSpec(body=MESSAGES_RESPONSE if route == "chat" else RESPONSES_RESPONSE)
- if route == "chat":
- kwargs: Final = {
- "model": MESSAGES_MODEL,
- "messages": list(MESSAGES),
- "api_key": "test-key",
- "api_base": server.base_url,
- "max_tokens": 32,
- **options,
- }
- request: Final = NativeCall(args=(), kwargs=kwargs, base={})
- return (_native.acompletion if asynchronous else _native.completion)(request)
- response_kwargs: Final = {
- "model": RESPONSES_MODEL,
- "input": "hello",
- "api_key": "test-key",
- "api_base": server.base_url,
- "max_output_tokens": 32,
- **options,
- }
- response_request: Final = NativeCall(
- args=(),
- kwargs=response_kwargs,
- base={
- "model": RESPONSES_MODEL,
- "input": "hello",
- "stream": None,
- "api_key": "test-key",
- "api_base": server.base_url,
- "custom_llm_provider": "openai",
- "extra_headers": None,
- },
- )
- return (_native.aresponses if asynchronous else _native.responses)(response_request)
-
-
-async def execute(route: Route, asynchronous: bool, server: RecordingServer, options: Mapping[str, object]) -> object:
- if not asynchronous:
- return await asyncio.to_thread(native_call, route, False, server, options)
- result: Final = native_call(route, True, server, options)
- assert isinstance(result, Awaitable)
- return await result
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_native_inference_returns_public_models_and_logs_once(
- route: Route,
- asynchronous: bool,
- recording_server: RecordingServer,
-) -> None:
- recorder: Final = RecordingLogger()
- result: Final = await execute(route, asynchronous, recording_server, {"callbacks": [recorder], "temperature": 0.25})
- assert len(recording_server.requests) == 1
- sent: Final = recording_server.requests[0]
- body: Final = _OBJECT.validate_python(sent.body)
- assert body["temperature"] == 0.25
- assert request_body(_OBJECT.validate_python(recorder.wait_for("log_pre_api_call")[0].kwargs)) == body
- if route == "chat":
- assert isinstance(result, ModelResponse)
- assert result.choices[0].message.content == "Hello from native Messages"
- assert sent.path == "/v1/messages"
- else:
- assert isinstance(result, ResponsesAPIResponse)
- assert result.output_text == "native response"
- assert sent.path == "/responses"
- success: Final = await recorder.wait_for_async("async_log_success_event" if asynchronous else "log_success_event")
- assert len(success) == 1
- if isinstance(result, ModelResponse):
- assert success[0].response is result
- else:
- logged: Final = success[0].response
- assert isinstance(logged, ResponsesAPIResponse)
- assert isinstance(result, ResponsesAPIResponse)
- assert logged.id == result.id
- assert logged.output_text == result.output_text
-
-
-@pytest.mark.asyncio
-async def test_native_inference_pre_call_edits_reach_the_provider(
- route: Route, recording_server: RecordingServer
-) -> None:
- class Edit(CustomLogger):
- def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
- request_body(kwargs)["temperature"] = 0.75
-
- await execute(route, True, recording_server, {"callbacks": [Edit()]})
- assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.75
-
-
-@pytest.mark.asyncio
-async def test_native_inference_does_not_restore_keywords_deleted_by_a_pre_call_hook(
- route: Route, recording_server: RecordingServer
-) -> None:
- class DropKeywords(CustomLogger):
- async def async_pre_call_deployment_hook(
- self, kwargs: dict[str, object], call_type: CallTypes | None
- ) -> dict[str, object]:
- return {name: value for name, value in kwargs.items() if name not in ("temperature", "extra_headers")}
-
- litellm.callbacks.append(DropKeywords())
- await execute(
- route, True, recording_server, {"temperature": 0.25, "extra_headers": {"x-deleted": "remove this header"}}
- )
-
- assert len(recording_server.requests) == 1
- sent: Final = recording_server.requests[0]
- assert "temperature" not in _OBJECT.validate_python(sent.body)
- assert sent.headers.get("x-deleted") is None
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("from_credentials", (False, True))
-async def test_native_resource_setup_uses_deployment_hook_arguments(
- route: Route,
- recording_server: RecordingServer,
- monkeypatch: pytest.MonkeyPatch,
- from_credentials: bool,
-) -> None:
- invalid_settings: Final = {"ssl_verify": object()}
- credential: Final = CredentialItem(
- credential_name="resource-settings", credential_info={}, credential_values=invalid_settings
- )
- monkeypatch.setattr(litellm, "credential_list", [credential])
-
- class Prepare(CustomLogger):
- async def async_pre_call_deployment_hook(
- self, kwargs: dict[str, object], call_type: CallTypes | None
- ) -> dict[str, object]:
- return {
- **kwargs,
- **({"litellm_credential_name": credential.credential_name} if from_credentials else invalid_settings),
- }
-
- litellm.callbacks.append(Prepare())
- recorder: Final = RecordingLogger()
- recording_server.expected_requests = 0
- with pytest.raises(ValueError, match=r"request\.ssl_verify") as caught:
- await execute(route, True, recording_server, {"callbacks": [recorder]})
- failure: Final = await recorder.wait_for_async("async_log_failure_event")
- assert len(failure) == 1
- assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
- assert not recording_server.requests
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_sdk_policy_rejection_precedes_resource_setup_and_is_logged_once(
- route: Route,
- asynchronous: bool,
- recording_server: RecordingServer,
- monkeypatch: pytest.MonkeyPatch,
-) -> None:
- monkeypatch.setattr(litellm, "max_budget", 1.0)
- monkeypatch.setattr(litellm, "_current_cost", 2.0)
- monkeypatch.setattr(litellm, "ssl_verify", object())
- recorder: Final = RecordingLogger()
- recording_server.expected_requests = 0
- with pytest.raises(litellm.BudgetExceededError) as caught:
- await execute(route, asynchronous, recording_server, {"callbacks": [recorder]})
- failure: Final = await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event")
- assert len(failure) == 1
- assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
- assert not recording_server.requests
- assert not any("success" in name for name in recorder.names)
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_native_inference_provider_failure_is_terminal_and_shared_with_callbacks(
- route: Route,
- asynchronous: bool,
- recording_server: RecordingServer,
-) -> None:
- recorder: Final = RecordingLogger()
- recording_server.enqueue(
- ResponseSpec(body={"error": {"message": "slow down", "type": "rate_limit_error"}}, status=429)
- )
- with pytest.raises(RateLimitError) as caught:
- await execute(route, asynchronous, recording_server, {"callbacks": [recorder]})
- assert getattr(caught.value, "status_code", None) == 429
- assert len(recording_server.requests) == 1
- failure: Final = await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event")
- assert len(failure) == 1
- assert _OBJECT.validate_python(failure[0].kwargs)["exception"] is caught.value
- assert not any("success" in name for name in recorder.names)
-
-
-@pytest.mark.asyncio
-async def test_unstarted_native_inference_has_no_provider_or_callback_effects(
- route: Route,
- recording_server: RecordingServer,
- monkeypatch: pytest.MonkeyPatch,
-) -> None:
- recorder: Final = RecordingLogger()
- recording_server.expected_requests = 0
- monkeypatch.setattr(litellm, "ssl_verify", object())
- monkeypatch.setattr(litellm, "max_budget", 1.0)
- monkeypatch.setattr(litellm, "_current_cost", 2.0)
- pending: Final = native_call(route, True, recording_server, {"callbacks": [recorder]})
- assert asyncio.iscoroutine(pending)
- pending.close()
- assert not recording_server.requests
- assert not recorder.events
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
- "options",
- (
- {"stream": True},
- {"extra_body": {"provider_option": True}},
- {"mock_response": "mock"},
- {"num_retries": 1},
- {"use_chat_completions_api": True},
- {"model_list": []},
- ),
-)
-async def test_native_responses_declines_unsupported_requests_before_callbacks(
- recording_server: RecordingServer,
- options: Mapping[str, object],
-) -> None:
- recorder: Final = RecordingLogger()
- recording_server.expected_requests = 0
- with pytest.raises(_native.RustBridgeDeclined):
- native_call("responses", True, recording_server, {**options, "callbacks": [recorder]})
- assert not recording_server.requests
- assert not recorder.events
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_native_chat_validation_failure_is_terminal(
- asynchronous: bool, recording_server: RecordingServer
-) -> None:
- recording_server.expected_requests = 0
- with pytest.raises(Exception, match="chat completions requires at least one message") as failure:
- await execute("chat", asynchronous, recording_server, {"messages": []})
- assert not isinstance(failure.value, _native.RustBridgeDeclined)
- assert not recording_server.requests
-
-
-@pytest.mark.asyncio
-async def test_native_projection_reads_positional_parameters(route: Route, recording_server: RecordingServer) -> None:
- from litellm.chat_completions.dispatch import (
- _DISPATCH as chat_dispatch, # pyright: ignore[reportPrivateUsage] # exercise the request passed to the native boundary
- )
- from litellm.responses.dispatch import (
- _DISPATCH as responses_dispatch, # pyright: ignore[reportPrivateUsage] # exercise the request passed to the native boundary
- )
-
- recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE if route == "chat" else RESPONSES_RESPONSE)
- kwargs: Final = {"api_key": "test-key", "api_base": recording_server.base_url}
- if route == "chat":
- args: Final = (MESSAGES_MODEL, list(MESSAGES), 12.0, 0.25)
- request: Final = chat_dispatch.request(args, kwargs)
- assert request is not None
- await asyncio.to_thread(_native.completion, request)
- assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.25
- else:
- response_args: Final = ("hello", RESPONSES_MODEL, None, "Be brief", 16)
- response_request: Final = responses_dispatch.request(response_args, kwargs)
- assert response_request is not None
- await asyncio.to_thread(_native.responses, response_request)
- body: Final = _OBJECT.validate_python(recording_server.requests[0].body)
- assert body["instructions"] == "Be brief"
- assert body["max_output_tokens"] == 16
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-@pytest.mark.parametrize("source", ("explicit", "base_url", "global", "provider", "environment", "empty"))
-async def test_native_connection_settings_reach_the_provider(
- route: Route,
- asynchronous: bool,
- source: str,
- recording_server: RecordingServer,
- monkeypatch: pytest.MonkeyPatch,
-) -> None:
- key: Final = "selected-key"
- explicit: Final = source in ("explicit", "base_url")
- monkeypatch.setattr(litellm, "api_key", key if source in ("global", "empty") else ("unused" if explicit else None))
- monkeypatch.setattr(litellm, "openai_key", key if source == "provider" else ("unused" if explicit else None))
- monkeypatch.setattr(litellm, "anthropic_key", key if source == "provider" else ("unused" if explicit else None))
- monkeypatch.setattr(
- litellm,
- "api_base",
- None if source == "environment" else ("http://127.0.0.1:1" if explicit else recording_server.base_url),
- )
- monkeypatch.setenv(
- "OPENAI_API_KEY" if route == "responses" else "ANTHROPIC_API_KEY", key if source == "environment" else "unused"
- )
- for name in ("OPENAI_BASE_URL", "OPENAI_API_BASE", "ANTHROPIC_BASE_URL", "ANTHROPIC_API_BASE"):
- monkeypatch.setenv(name, recording_server.base_url if source == "environment" else "http://127.0.0.1:1")
- result: Final = await execute(
- route,
- asynchronous,
- recording_server,
- {
- "api_key": key if explicit else ("" if source == "empty" else None),
- "api_base": recording_server.base_url if source == "explicit" else ("" if source == "empty" else None),
- **({"base_url": recording_server.base_url} if source == "base_url" else {}),
- },
- )
- assert isinstance(result, ModelResponse | ResponsesAPIResponse)
- assert len(recording_server.requests) == 1
- headers: Final = recording_server.requests[0].headers
- assert headers["x-api-key" if route == "chat" else "authorization"] == (key if route == "chat" else f"Bearer {key}")
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-@pytest.mark.parametrize("encoded", (False, True))
-async def test_native_responses_decode_continuation_ids(
- asynchronous: bool, encoded: bool, recording_server: RecordingServer
-) -> None:
- original: Final = "resp_upstream"
- previous: Final = (
- ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "deployment", original)
- if encoded
- else original
- )
- await execute("responses", asynchronous, recording_server, {"previous_response_id": previous})
- assert _OBJECT.validate_python(recording_server.requests[0].body)["previous_response_id"] == original
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("asynchronous", (False, True))
-async def test_native_chat_uses_bound_positional_parameters(
- asynchronous: bool, recording_server: RecordingServer
-) -> None:
- recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
- arguments: Final = (MESSAGES_MODEL, list(MESSAGES), 12.0, 0.35)
- supplied: Final = {"base_url": recording_server.base_url, "api_key": "test-key", "max_tokens": 32}
- call: Final = chat_dispatch._DISPATCH.request(arguments, supplied) # pyright: ignore[reportPrivateUsage] # exercise the native request produced by public binding
- assert call is not None
- result: Final = (
- await _native.acompletion(call) if asynchronous else await asyncio.to_thread(_native.completion, call)
- )
- assert isinstance(result, ModelResponse)
- assert result.choices[0].message.content == "Hello from native Messages"
- assert len(recording_server.requests) == 1
- body: Final = _OBJECT.validate_python(recording_server.requests[0].body)
- assert body["temperature"] == arguments[3]
- assert body["max_tokens"] == supplied["max_tokens"]
diff --git a/tests/unit/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py
index 89fd1d1b771..64208182025 100644
--- a/tests/unit/rust_bridge/messages/test_secrets.py
+++ b/tests/unit/rust_bridge/messages/test_secrets.py
@@ -14,8 +14,8 @@ from litellm.rust_bridge.host import settings
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES
from litellm.rust_bridge.public_call import NativeCall
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
+from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service
-from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE
pytest.importorskip("litellm.rust_bridge._native")
diff --git a/tests/unit/rust_bridge/ocr/test_secrets.py b/tests/unit/rust_bridge/ocr/test_secrets.py
index a9c6072b357..33ded916888 100644
--- a/tests/unit/rust_bridge/ocr/test_secrets.py
+++ b/tests/unit/rust_bridge/ocr/test_secrets.py
@@ -16,8 +16,8 @@ from litellm.rust_bridge.host import settings
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
from litellm.rust_bridge.public_call import NativeCall
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
+from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service
-from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE
native: Final = pytest.importorskip("litellm.rust_bridge._native")
From 327bdcaa7f978719df403e4c8a48dea8fd6bb672 Mon Sep 17 00:00:00 2001
From: yujonglee
Date: Wed, 14 Oct 2026 17:19:05 -0700
Subject: [PATCH 061/140] test(host-python): move driver tests out of the
inline module (#46325)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm-rust/crates/host-python/src/driver.rs | 1840 -----------------
.../crates/host-python/tests/driver.rs | 866 ++++++++
.../crates/host-python/tests/handle.rs | 201 ++
.../crates/host-python/tests/support/mod.rs | 656 ++++++
4 files changed, 1723 insertions(+), 1840 deletions(-)
create mode 100644 litellm-rust/crates/host-python/tests/driver.rs
create mode 100644 litellm-rust/crates/host-python/tests/handle.rs
create mode 100644 litellm-rust/crates/host-python/tests/support/mod.rs
diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs
index c70f2a5f1a0..a1450ed6308 100644
--- a/litellm-rust/crates/host-python/src/driver.rs
+++ b/litellm-rust/crates/host-python/src/driver.rs
@@ -701,1843 +701,3 @@ where
self.clear();
}
}
-
-#[cfg(test)]
-mod tests {
- use std::sync::{Arc, Mutex};
-
- use litellm_host::{
- interceptors::{RawResponse, RequestContext},
- machine::{CallMachine, MachineFault},
- };
- use pyo3::exceptions::{PyBaseException, PyValueError};
- use pyo3::types::PyDict;
-
- use super::*;
- use crate::{PythonOwned, PythonRuntime};
- use litellm_host::hooks::CallHooks;
- use litellm_host::interceptors::Interceptors;
-
- static PYTHON_GLOBALS: Mutex<()> = Mutex::new(());
-
- fn lifecycle_binding(py: Python<'_>) -> PyResult> {
- py.import("host_test_lifecycle")
- }
-
- fn call_options(asynchronous: bool) -> CallOptions {
- CallOptions::new(asynchronous, lifecycle_binding)
- }
-
- fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
- let source =
- std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/lifecycle.py"))
- .unwrap();
- PyModule::from_code(
- py,
- &source,
- pyo3::ffi::c_str!("lifecycle.py"),
- pyo3::ffi::c_str!("host_test_lifecycle"),
- )
- .unwrap()
- }
-
- #[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)
- }
- }
-
- impl From for Error {
- fn from(fault: MachineFault) -> Self {
- Self(format!("{fault:?}"))
- }
- }
-
- struct Synthetic;
-
- impl Protocol for Synthetic {
- type Response = String;
- type Error = Error;
- type Request = String;
- type HostCall = (&'static str, Reply);
- type Chunk = std::convert::Infallible;
- type StreamHead = std::convert::Infallible;
- }
-
- fn wire() -> WireRequest {
- WireRequest {
- url: "https://example.invalid".into(),
- headers: Vec::new(),
- body: serde_json::json!({}),
- }
- }
-
- fn context() -> RequestContext {
- RequestContext {
- model: "model".into(),
- custom_llm_provider: "provider".into(),
- optional_params: serde_json::json!({}),
- secret_fields: Vec::new(),
- api_key: None,
- }
- }
-
- #[derive(Default)]
- struct Log(Arc>>);
-
- impl Log {
- fn push(&self, entry: impl Into) {
- self.0.lock().unwrap().push(entry.into());
- }
-
- fn entries(&self) -> Vec {
- self.0.lock().unwrap().clone()
- }
- }
-
- #[derive(Clone, Copy)]
- enum OpScript {
- Answer,
- RaisePython,
- RejectNatively,
- RejectRequestNatively,
- RaiseRequestPython,
- AwaitAnswer,
- AwaitFailure,
- }
-
- struct SyntheticBinding {
- log: Log,
- op: OpScript,
- classifier_fails: bool,
- pending_reply: Option>,
- }
-
- /// 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 SyntheticBinding {
- fn answer(&self, value: impl FnOnce() -> String) -> Result> {
- match self.op {
- OpScript::Answer | OpScript::AwaitAnswer | OpScript::AwaitFailure => Ok(value()),
- OpScript::RaisePython | OpScript::RaiseRequestPython => {
- Err(PyValueError::new_err("op failed").into())
- }
- OpScript::RejectNatively | OpScript::RejectRequestNatively => {
- Err(InvokeError::Native(Error("op rejected".into())))
- }
- }
- }
- }
-
- impl PythonBinding for SyntheticBinding {
- type Protocol = Synthetic;
- type Failure = Classified;
-
- fn decode_request(
- &mut self,
- _: Python<'_>,
- arguments: &Bound<'_, PyDict>,
- ) -> Result> {
- self.log.push("project");
- match self.op {
- OpScript::RejectRequestNatively | OpScript::RaiseRequestPython => {
- self.answer(String::new)
- }
- _ => Ok(format!("project:{}", arguments.len())),
- }
- }
-
- fn encode_stream_head(
- &mut self,
- _: Python<'_>,
- head: std::convert::Infallible,
- ) -> PyResult> {
- match head {}
- }
-
- fn encode_chunk(
- &mut self,
- _: Python<'_>,
- chunk: std::convert::Infallible,
- ) -> PyResult> {
- match chunk {}
- }
-
- fn encode_response(&mut self, py: Python<'_>, response: String) -> PyResult> {
- self.log.push("complete");
- Ok(pyo3::types::PyString::new(py, &response)
- .into_any()
- .unbind())
- }
-
- fn map_error(&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())
- }
- }
-
- impl PythonHostCalls for SyntheticBinding {
- fn handle_host_call(
- &mut self,
- _: Python<'_>,
- (op, reply): (&'static str, Reply),
- ) -> Result<(), InvokeError> {
- self.log.push(format!("op:{op}"));
- self.answer(|| op.to_string())
- .map(|answer| reply.send(answer))
- }
-
- fn begin_host_call(
- &mut self,
- py: Python<'_>,
- (op, reply): (&'static str, Reply),
- ) -> Result>, InvokeError> {
- if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) {
- return self.handle_host_call(py, (op, reply)).map(|()| None);
- }
- self.pending_reply = Some(reply);
- let module = PyModule::from_code(
- py,
- pyo3::ffi::c_str!(
- "async def answer(fail):\n if fail:\n raise LookupError('async host failed')\n return 'awaited'\n"
- ),
- pyo3::ffi::c_str!("host_op.py"),
- pyo3::ffi::c_str!("host_op"),
- )?;
- Ok(Some(
- module
- .getattr("answer")?
- .call1((matches!(self.op, OpScript::AwaitFailure),))?
- .unbind(),
- ))
- }
-
- fn resume_host_call(
- &mut self,
- py: Python<'_>,
- result: PyResult>,
- ) -> Result>, InvokeError> {
- let answer = result?.extract::(py)?;
- self.pending_reply.take().unwrap().send(answer);
- Ok(None)
- }
- }
-
- impl PythonOwned for SyntheticBinding {
- fn close(&mut self, _: Python<'_>) {
- self.log.push("host.close");
- }
- fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
- Ok(())
- }
- }
-
- #[derive(Clone, Copy)]
- enum HookScript {
- Plain,
- RewriteArguments,
- ObserveArguments,
- FailBegin,
- ReplaceResponse,
- FailAfterSuccess,
- CancelTerminal,
- FailTerminal,
- }
-
- struct SyntheticHooks {
- log: Log,
- script: HookScript,
- }
-
- impl CallHooks for SyntheticHooks {
- fn prepare_arguments(
- &mut self,
- _: Python<'_>,
- arguments: Py,
- _: f64,
- ) -> PyResult>> {
- self.log.push("begin");
- if matches!(self.script, HookScript::FailBegin) {
- return Err(PyValueError::new_err("begin failed"));
- }
- if matches!(self.script, HookScript::RewriteArguments) {
- let prepared = Python::attach(|py| -> PyResult> {
- let prepared = arguments.bind(py).copy()?;
- prepared.set_item("prepared", "hook")?;
- prepared.set_item("api_key", "hook-key")?;
- Ok(prepared.unbind())
- })?;
- return Ok(HookStep::Ready(prepared));
- }
- Ok(HookStep::Ready(arguments))
- }
-
- fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> {
- if matches!(self.script, HookScript::ObserveArguments) {
- let value: String = arguments
- .bind(py)
- .get_item("prepared")?
- .unwrap()
- .extract()?;
- self.log.push(format!("adopted:{value}"));
- }
- Ok(())
- }
-
- fn before_provider_request(
- &mut self,
- _: Python<'_>,
- wire: Box,
- _: &RequestContext,
- ) -> PyResult>> {
- self.log.push("before_provider_request");
- Ok(HookStep::Ready(Box::new(WireRequest {
- url: "rewritten".into(),
- ..*wire
- })))
- }
-
- fn transform_response(
- &mut self,
- py: Python<'_>,
- response: Py,
- _: Timing,
- ) -> PyResult>> {
- self.log.push("after_success");
- match self.script {
- HookScript::ReplaceResponse => Ok(HookStep::Ready(
- "replaced".into_pyobject(py)?.into_any().unbind(),
- )),
- HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")),
- HookScript::Plain
- | HookScript::ObserveArguments
- | HookScript::RewriteArguments
- | HookScript::FailBegin
- | HookScript::CancelTerminal
- | HookScript::FailTerminal => Ok(HookStep::Ready(response)),
- }
- }
-
- fn on_event(
- &mut self,
- py: Python<'_>,
- event: PythonCallEvent<'_>,
- ) -> PyResult> {
- if matches!(self.script, HookScript::CancelTerminal)
- && matches!(
- event,
- PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
- )
- {
- return Err(pyo3::exceptions::asyncio::CancelledError::new_err(
- "callback cancelled",
- ));
- }
- if matches!(self.script, HookScript::FailTerminal)
- && matches!(
- event,
- PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
- )
- {
- return Err(PyValueError::new_err("callback failed"));
- }
- self.log.push(match event {
- PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => {
- "cache_hit".into()
- }
- PythonCallEvent::Started { .. } => "started".into(),
- PythonCallEvent::Cancelled { .. } => "cancelled".into(),
- PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
- format!("response:{}", raw.body)
- }
- PythonCallEvent::Succeeded { response, .. } => {
- format!("succeeded:{}", response.bind(py))
- }
- PythonCallEvent::Failed { origin, error, .. } => {
- format!("failed:{origin:?}:{}", error.value(py))
- }
- });
- Ok(HookStep::Ready(()))
- }
-
- fn on_stream_open(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> {
- self.log.push("opened");
- Ok(())
- }
-
- fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> {
- self.log.push("delivered");
- Ok(())
- }
- }
-
- impl PythonOwned for SyntheticHooks {
- fn close(&mut self, _: Python<'_>) {
- self.log.push("adapter.close");
- }
- fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
- Ok(())
- }
- }
-
- fn run_scripted(
- py: Python<'_>,
- machine: impl FnOnce(String) -> CallMachine + Send + Sync + 'static,
- op: OpScript,
- script: HookScript,
- asynchronous: bool,
- ) -> (PyResult>, Vec) {
- run_hosted(
- py,
- machine,
- SyntheticBinding {
- log: Log::default(),
- op,
- classifier_fails: false,
- pending_reply: None,
- },
- script,
- asynchronous,
- )
- }
-
- fn run_hosted(
- py: Python<'_>,
- machine: impl FnOnce(String) -> CallMachine + Send + Sync + 'static,
- host: SyntheticBinding,
- script: HookScript,
- asynchronous: bool,
- ) -> (PyResult>, Vec) {
- run_composed(
- py,
- machine,
- host,
- script,
- std::convert::identity,
- call_options(asynchronous),
- )
- }
-
- #[rstest::rstest]
- #[case::success(OpScript::AwaitAnswer)]
- #[case::failure(OpScript::AwaitFailure)]
- fn asynchronous_host_operations_resume_the_same_machine(#[case] op: OpScript) {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(std::sync::PoisonError::into_inner);
- Python::initialize();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let (result, log) = run_scripted(
- py,
- |_| {
- CallMachine::::new(None, |host| {
- Box::pin(async move { host.services.call(|reply| ("read", reply)).await })
- })
- },
- op,
- HookScript::Plain,
- true,
- );
- match op {
- OpScript::AwaitAnswer => {
- assert_eq!(result.unwrap().extract::(py).unwrap(), "awaited")
- }
- OpScript::AwaitFailure => {
- assert!(
- result
- .unwrap_err()
- .is_instance_of::(py)
- );
- assert_eq!(
- log.iter()
- .filter(|entry| entry.starts_with("failed:"))
- .count(),
- 1
- );
- assert!(!log.iter().any(|entry| entry.starts_with("succeeded:")));
- }
- _ => unreachable!(),
- }
- });
- }
-
- #[rstest::rstest]
- #[case::synchronous(false)]
- #[case::asynchronous(true)]
- fn composed_hooks_share_prepared_arguments_and_one_execution(#[case] asynchronous: bool) {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let log = Log::default();
- let hook = |script| SyntheticHooks {
- log: Log(log.0.clone()),
- script,
- };
- let hooks = crate::HookChain::new()
- .with(hook(HookScript::ObserveArguments))
- .with(hook(HookScript::RewriteArguments))
- .with(hook(HookScript::ReplaceResponse));
- let result = run_call(
- py,
- move |_, _, request| Ok(success_machine()(request)),
- SyntheticBinding {
- log: Log(log.0.clone()),
- op: OpScript::Answer,
- classifier_fails: false,
- pending_reply: None,
- },
- hooks,
- PyDict::new(py).unbind(),
- call_options(asynchronous),
- )
- .unwrap();
- let response = if asynchronous {
- assert!(log.entries().is_empty());
- let completion = result.call_method1(py, "send", (py.None(),)).unwrap_err();
- assert!(completion.is_instance_of::(py));
- completion.value(py).getattr("value").unwrap().unbind()
- } else {
- result
- };
- assert_eq!(response.extract::(py).unwrap(), "replaced");
- let entries = log.entries();
- assert!(entries.iter().any(|entry| entry == "adopted:hook"));
- assert_eq!(
- entries
- .iter()
- .filter(|entry| entry.as_str() == "project")
- .count(),
- 1
- );
- assert_eq!(
- entries
- .iter()
- .filter(|entry| entry.as_str() == "op:sign")
- .count(),
- 1
- );
- assert_eq!(
- entries
- .iter()
- .filter(|entry| entry.as_str() == "succeeded:replaced")
- .count(),
- 3
- );
- assert!(!entries.iter().any(|entry| entry.starts_with("failed:")));
- });
- }
-
- fn run_composed(
- py: Python<'_>,
- machine: impl FnOnce(String) -> CallMachine + Send + Sync + 'static,
- host: SyntheticBinding,
- script: HookScript,
- compose: impl FnOnce(SyntheticHooks) -> L,
- options: CallOptions,
- ) -> (PyResult>, Vec) {
- let asynchronous = options.asynchronous;
- let log = Log(host.log.0.clone());
- let adapter = SyntheticHooks {
- log: Log(log.0.clone()),
- script,
- };
- let arguments = PyDict::new(py);
- arguments.set_item("model", "m").unwrap();
- let result = run_call(
- py,
- move |_, _, request| Ok(machine(request)),
- host,
- compose(adapter),
- arguments.unbind(),
- options,
- );
- let result = if asynchronous {
- result.and_then(|coroutine| {
- let completed = coroutine
- .call_method1(py, "send", (py.None(),))
- .unwrap_err();
- if !completed.is_instance_of::(py) {
- return Err(completed);
- }
- completed.value(py).getattr("value").map(Bound::unbind)
- })
- } else {
- result
- };
- (result, log.entries())
- }
-
- /// Answers to projection, to the route op and to `before_provider_request` all reach the
- /// response, so a driver that misroutes a reply changes what the call returns.
- fn success_machine() -> impl FnOnce(String) -> CallMachine + Send + Sync {
- move |projected| {
- CallMachine::new(None, move |host| {
- Box::pin(async move {
- let signed = host.services.call(|reply| ("sign", reply)).await?;
- let wire = host
- .interceptors
- .before_provider_request(wire(), context())
- .await?;
- host.interceptors
- .after_provider_response(RawResponse { body: "raw".into() })
- .await?;
- Ok(format!("{projected}|{signed}|{}", wire.url))
- })
- })
- }
- }
-
- #[rstest::rstest]
- #[case::draining(4, false)]
- #[case::full(1, false)]
- #[case::closed(4, true)]
- fn observation_delivery_preserves_callback_order_and_response(
- #[case] capacity: usize,
- #[case] closed: bool,
- #[values(false, true)] asynchronous: bool,
- ) {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let (expected, expected_log) = run_scripted(
- py,
- success_machine(),
- OpScript::Answer,
- HookScript::ReplaceResponse,
- asynchronous,
- );
- let (sender, mut receiver) = litellm_host::observation::observation_channel(
- std::num::NonZeroUsize::new(capacity).unwrap(),
- );
- if closed {
- receiver.close();
- }
- let (actual, actual_log) = run_composed(
- py,
- success_machine(),
- SyntheticBinding {
- log: Log::default(),
- op: OpScript::Answer,
- classifier_fails: false,
- pending_reply: None,
- },
- HookScript::ReplaceResponse,
- std::convert::identity,
- CallOptions {
- asynchronous,
- lifecycle: lifecycle_binding,
- observers: Some(sender.clone()),
- },
- );
- assert_eq!(
- actual.unwrap().extract::(py).unwrap(),
- expected.unwrap().extract::(py).unwrap()
- );
- assert_eq!(actual_log, expected_log);
- let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
- assert_eq!(events.len() as u64 + sender.dropped_events(), 3);
- if !closed && capacity >= 3 {
- let [
- CallEvent::Started { start_time },
- CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
- CallEvent::Succeeded { timing, .. },
- ] = events.as_slice()
- else {
- panic!("expected start, provider response and success: {events:?}");
- };
- assert_eq!(raw.body, "raw");
- assert_eq!(*start_time, timing.start_time);
- assert!(timing.end_time >= timing.start_time);
- }
- });
- }
-
- #[rstest::rstest]
- #[case::preparation(HookScript::FailBegin, false, 1)]
- #[case::transformation(HookScript::FailAfterSuccess, false, 1)]
- #[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
- #[case::terminal_failure(HookScript::FailTerminal, false, 0)]
- fn observation_reports_the_final_outcome_without_replaying_callbacks(
- #[case] script: HookScript,
- #[case] cancelled: bool,
- #[case] failure_callbacks: usize,
- #[values(false, true)] asynchronous: bool,
- ) {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let (sender, mut receiver) = litellm_host::observation::observation_channel(
- std::num::NonZeroUsize::new(4).unwrap(),
- );
- let (result, log) = run_composed(
- py,
- success_machine(),
- SyntheticBinding {
- log: Log::default(),
- op: OpScript::Answer,
- classifier_fails: false,
- pending_reply: None,
- },
- script,
- std::convert::identity,
- CallOptions {
- asynchronous,
- lifecycle: lifecycle_binding,
- observers: Some(sender),
- },
- );
- let error = result.unwrap_err();
- assert_eq!(
- error.is_instance_of::(py),
- cancelled
- );
- let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
- assert!(matches!(events.first(), Some(CallEvent::Started { .. })));
- assert_eq!(
- matches!(events.last(), Some(CallEvent::Cancelled { .. })),
- cancelled
- );
- assert_eq!(
- matches!(
- events.last(),
- Some(CallEvent::Failed {
- origin: FailureOrigin::Host,
- ..
- })
- ),
- !cancelled
- );
- assert_eq!(
- events
- .iter()
- .filter(|event| matches!(
- event,
- CallEvent::Failed { .. }
- | CallEvent::Succeeded { .. }
- | CallEvent::Cancelled { .. }
- ))
- .count(),
- 1
- );
- assert_eq!(
- log.iter()
- .filter(|entry| entry.starts_with("failed:"))
- .count(),
- failure_callbacks
- );
- assert!(
- !log.iter()
- .any(|entry| entry.starts_with("succeeded:") || entry == "cancelled")
- );
- });
- }
-
- #[rstest::rstest]
- #[case::preparation(HookScript::FailBegin, OpScript::Answer)]
- #[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
- #[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)]
- fn startup_failures_release_the_factory_without_constructing_a_machine(
- #[case] script: HookScript,
- #[case] op: OpScript,
- #[values(false, true)] asynchronous: bool,
- ) {
- use std::sync::atomic::{AtomicBool, Ordering};
- struct Release(Arc);
- impl Drop for Release {
- fn drop(&mut self) {
- self.0.store(true, Ordering::SeqCst);
- }
- }
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let constructed = Arc::new(AtomicBool::new(false));
- let did_construct = constructed.clone();
- let released = Arc::new(AtomicBool::new(false));
- let release = Release(released.clone());
- let (result, log) = run_scripted(
- py,
- move |request| {
- let _release = release;
- did_construct.store(true, Ordering::SeqCst);
- success_machine()(request)
- },
- op,
- script,
- asynchronous,
- );
- assert!(result.is_err());
- assert!(!constructed.load(Ordering::SeqCst));
- assert!(released.load(Ordering::SeqCst));
- assert_eq!(
- log.iter()
- .filter(|event| event.starts_with("failed:"))
- .count(),
- 1
- );
- assert_eq!(
- log.iter().filter(|event| *event == "adapter.close").count(),
- 1
- );
- assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1);
- });
- }
-
- #[rstest::rstest]
- fn success_runs_every_step_in_order_and_returns_the_public_response() {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- for asynchronous in [false, true] {
- let (result, log) = run_scripted(
- py,
- success_machine(),
- OpScript::Answer,
- HookScript::Plain,
- asynchronous,
- );
- assert_eq!(
- result.unwrap().extract::(py).unwrap(),
- "project:1|sign|rewritten"
- );
- assert_eq!(
- log,
- [
- "started",
- "begin",
- "project",
- "op:sign",
- "before_provider_request",
- "response:raw",
- "complete",
- "after_success",
- "succeeded:project:1|sign|rewritten",
- "adapter.close",
- "host.close",
- ]
- );
- }
- });
- }
-
- struct Streaming;
-
- impl Protocol for Streaming {
- type Response = ();
- type Error = Error;
- type Request = ();
- type HostCall = std::convert::Infallible;
- type Chunk = &'static str;
- type StreamHead = Vec<(&'static str, &'static str)>;
- }
-
- struct StreamingBinding;
-
- impl PythonBinding for StreamingBinding {
- type Protocol = Streaming;
- type Failure = Classified;
-
- fn decode_request(
- &mut self,
- _: Python<'_>,
- _: &Bound<'_, PyDict>,
- ) -> Result<(), InvokeError> {
- Ok(())
- }
-
- fn encode_stream_head(
- &mut self,
- py: Python<'_>,
- head: Vec<(&'static str, &'static str)>,
- ) -> PyResult> {
- let headers = PyDict::new(py);
- for (name, value) in head {
- headers.set_item(name, value)?;
- }
- let hidden = PyDict::new(py);
- hidden.set_item("additional_headers", headers)?;
- Ok(hidden.into_any().unbind())
- }
-
- fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult> {
- Ok(pyo3::types::PyString::new(py, chunk).into_any().unbind())
- }
-
- fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult> {
- Ok(py.None())
- }
-
- fn map_error(&self, _: Python<'_>, error: Error) -> PyResult {
- Ok(Classified(error.0))
- }
-
- fn host_error(error: &PyErr) -> Error {
- Error(error.to_string())
- }
- }
-
- impl PythonHostCalls for StreamingBinding {
- fn handle_host_call(
- &mut self,
- _: Python<'_>,
- op: std::convert::Infallible,
- ) -> Result<(), InvokeError> {
- match op {}
- }
- }
-
- impl PythonOwned for StreamingBinding {
- fn close(&mut self, _: Python<'_>) {}
- fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
- Ok(())
- }
- }
-
- fn streaming_machine()
- -> impl FnOnce(()) -> litellm_host::call::HostedMachine + Send + Sync {
- move |()| {
- litellm_host::call::hosted_call((), None, |(), _, _, _observations| async {
- Ok(litellm_host::call::CallOutput::Stream {
- head: vec![("request-id", "req_1")],
- chunks: Box::pin(futures_util::stream::iter([Ok("first"), Ok("second")])),
- })
- })
- }
- }
-
- #[rstest::rstest]
- #[case::sync(false)]
- #[case::asynchronous(true)]
- fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
- #[case] asynchronous: bool,
- ) {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let log = Log::default();
- let (sender, mut receiver) = litellm_host::observation::observation_channel(
- std::num::NonZeroUsize::new(3).unwrap(),
- );
- let handed = run_call(
- py,
- |_, _, request| Ok(streaming_machine()(request)),
- StreamingBinding,
- SyntheticHooks {
- log: Log(log.0.clone()),
- script: HookScript::Plain,
- },
- PyDict::new(py).unbind(),
- CallOptions {
- asynchronous,
- lifecycle: lifecycle_binding,
- observers: Some(sender),
- },
- )
- .unwrap();
- let stream = if asynchronous {
- let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
- stop.value(py).getattr("value").unwrap()
- } else {
- handed.into_bound(py)
- };
- if asynchronous {
- for method in ["__anext__", "aclose", "aclose"] {
- let stop = stream
- .call_method0(method)
- .unwrap()
- .call_method1("send", (py.None(),))
- .unwrap_err();
- assert!(stop.is_instance_of::(py));
- }
- } else {
- assert_eq!(
- stream
- .call_method0("__next__")
- .unwrap()
- .extract::()
- .unwrap(),
- "first"
- );
- stream.call_method0("close").unwrap();
- stream.call_method0("close").unwrap();
- }
- let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
- assert!(matches!(
- events.as_slice(),
- [CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
- ));
- assert_eq!(
- log.entries(),
- [
- "started",
- "begin",
- "opened",
- "delivered",
- "succeeded:None",
- "adapter.close"
- ]
- );
- });
- }
-
- /// Drives a `Stream` (async) or `SyncStream` to completion from a sync test.
- fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec {
- if !asynchronous {
- return stream
- .try_iter()
- .unwrap()
- .map(|chunk| chunk.unwrap().extract().unwrap())
- .collect();
- }
- std::iter::from_fn(|| {
- let stop = stream
- .call_method0("__anext__")
- .unwrap()
- .call_method1("send", (py.None(),))
- .unwrap_err();
- if stop.is_instance_of::(py) {
- return None;
- }
- assert!(stop.is_instance_of::(py));
- Some(stop.value(py).getattr("value").unwrap().extract().unwrap())
- })
- .collect()
- }
-
- #[rstest::rstest]
- fn a_stream_carries_its_head_before_the_first_chunk() {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- for asynchronous in [false, true] {
- let log = Log::default();
- let adapter = SyntheticHooks {
- log: Log(log.0.clone()),
- script: HookScript::Plain,
- };
- let handed = run_call(
- py,
- |_, _, request| Ok(streaming_machine()(request)),
- StreamingBinding,
- adapter,
- PyDict::new(py).unbind(),
- call_options(asynchronous),
- )
- .unwrap();
- let stream = if asynchronous {
- let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
- stop.value(py).getattr("value").unwrap()
- } else {
- handed.into_bound(py)
- };
- let hidden: std::collections::HashMap<
- String,
- std::collections::HashMap,
- > = stream.getattr("head").unwrap().extract().unwrap();
- assert_eq!(
- hidden["additional_headers"],
- std::collections::HashMap::from([(
- "request-id".to_string(),
- "req_1".to_string()
- )])
- );
- assert_eq!(log.entries(), ["started", "begin", "opened"]);
- assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
- }
- });
- }
-
- fn failing_machine() -> impl FnOnce(String) -> CallMachine + Send + Sync {
- move |_| {
- CallMachine::new(None, |_| {
- Box::pin(async move { Err(Error("provider exploded".into())) })
- })
- }
- }
-
- #[rstest::rstest]
- 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| {
- install_lifecycle_module(py);
- for asynchronous in [false, true] {
- let (result, log) = run_scripted(
- py,
- failing_machine(),
- OpScript::Answer,
- HookScript::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",
- "project",
- "classify:provider exploded",
- "failed:Call:classified: provider exploded",
- "adapter.close",
- "host.close",
- ]
- );
- }
- });
- }
-
- #[rstest::rstest]
- 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,
- HookScript::Plain,
- false,
- );
- assert_eq!(
- result.unwrap_err().value(py).to_string(),
- "classified: op rejected"
- );
- assert_eq!(
- log,
- [
- "started",
- "begin",
- "project",
- "op:sign",
- "classify:op rejected",
- "failed:Call:classified: op rejected",
- "adapter.close",
- "host.close",
- ]
- );
- });
- }
-
- #[rstest::rstest]
- 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(),
- OpScript::RaisePython,
- HookScript::Plain,
- false,
- );
- let error = result.unwrap_err();
- assert!(error.is_instance_of::(py));
- assert_eq!(error.value(py).to_string(), "op failed");
- assert_eq!(
- log,
- [
- "started",
- "begin",
- "project",
- "op:sign",
- "failed:Call:op failed",
- "adapter.close",
- "host.close",
- ]
- );
- });
- }
-
- #[rstest::rstest]
- 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(),
- SyntheticBinding {
- log: Log::default(),
- op: OpScript::Answer,
- classifier_fails: true,
- pending_reply: None,
- },
- HookScript::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.context(py).unwrap();
- assert!(context.is_instance_of::(py));
- assert_eq!(context.value(py).to_string(), "provider exploded");
- assert_eq!(
- log,
- [
- "started",
- "begin",
- "project",
- "classify:provider exploded",
- "failed:Call:classifier failed",
- "adapter.close",
- "host.close",
- ]
- );
- });
- }
-
- #[rstest::rstest]
- fn begin_failures_are_host_failures_without_provider_mapping() {
- 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::Answer,
- HookScript::FailBegin,
- false,
- );
- let error = result.unwrap_err();
- assert_eq!(error.value(py).to_string(), "begin failed");
- assert_eq!(
- log,
- [
- "started",
- "begin",
- "failed:Host:begin failed",
- "adapter.close",
- "host.close"
- ]
- );
- });
- }
-
- enum ArgumentPolicy {
- Inherit,
- Reject(Py),
- }
-
- impl CallHooks for ArgumentPolicy {
- fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> {
- match self {
- Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
- Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
- }
- }
- }
-
- impl PythonOwned for ArgumentPolicy {
- fn close(&mut self, _: Python<'_>) {}
-
- fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
- match self {
- Self::Inherit => Ok(()),
- Self::Reject(error) => visit.call(error),
- }
- }
- }
-
- #[rstest::rstest]
- #[case::sync_success(false, false)]
- #[case::async_success(true, false)]
- #[case::sync_failure(false, true)]
- #[case::async_failure(true, true)]
- fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
- #[case] asynchronous: bool,
- #[case] fail_setup: bool,
- ) {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let log = Log::default();
- let setup_log = Log(log.0.clone());
- let error = PyValueError::new_err("setup failed").into_value(py);
- let setup_error = error.clone_ref(py);
- let arguments = PyDict::new(py);
- arguments.set_item("api_key", "original").unwrap();
- let result = run_call(
- py,
- move |py, prepared, request| {
- let source: String = prepared.get_item("prepared")?.unwrap().extract()?;
- let key: String = prepared.get_item("api_key")?.unwrap().extract()?;
- setup_log.push(format!("setup:{source}:{key}"));
- if fail_setup {
- return Err(PyErr::from_value(setup_error.into_bound(py).into_any()));
- }
- Ok(success_machine()(request))
- },
- SyntheticBinding {
- log: Log(log.0.clone()),
- op: OpScript::Answer,
- classifier_fails: false,
- pending_reply: None,
- },
- crate::HookChain::new()
- .with(SyntheticHooks {
- log: Log(log.0.clone()),
- script: HookScript::RewriteArguments,
- })
- .with(ArgumentPolicy::Inherit),
- arguments.clone().unbind(),
- call_options(asynchronous),
- );
- let settled = if asynchronous {
- assert!(log.entries().is_empty());
- let completion = result
- .unwrap()
- .call_method1(py, "send", (py.None(),))
- .unwrap_err();
- if completion.is_instance_of::(py) {
- completion.value(py).getattr("value").map(Bound::unbind)
- } else {
- Err(completion)
- }
- } else {
- result
- };
- assert_eq!(
- arguments
- .get_item("api_key")
- .unwrap()
- .unwrap()
- .extract::()
- .unwrap(),
- "original"
- );
- assert_eq!(
- &log.entries()[..4],
- ["started", "begin", "project", "setup:hook:inherited"]
- );
- if fail_setup {
- assert!(settled.unwrap_err().value(py).is(error.bind(py)));
- assert_eq!(
- log.entries(),
- [
- "started",
- "begin",
- "project",
- "setup:hook:inherited",
- "failed:Host:setup failed",
- "adapter.close",
- "host.close",
- ]
- );
- } else {
- assert!(settled.is_ok());
- assert!(
- log.entries()
- .iter()
- .any(|entry| entry.starts_with("succeeded:"))
- );
- }
- });
- }
-
- #[rstest::rstest]
- fn closing_an_unstarted_call_never_initializes_resources() {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let log = Log::default();
- let setup_log = Log(log.0.clone());
- let pending = run_call(
- py,
- move |_, _, request| {
- setup_log.push("setup");
- Ok(success_machine()(request))
- },
- SyntheticBinding {
- log: Log(log.0.clone()),
- op: OpScript::Answer,
- classifier_fails: false,
- pending_reply: None,
- },
- SyntheticHooks {
- log: Log(log.0.clone()),
- script: HookScript::Plain,
- },
- PyDict::new(py).unbind(),
- call_options(true),
- )
- .unwrap();
- assert!(log.entries().is_empty());
- pending.call_method0(py, "close").unwrap();
- drop(pending);
- assert_eq!(log.entries(), ["adapter.close", "host.close"]);
- });
- }
-
- #[rstest::rstest]
- #[case::synchronous(false)]
- #[case::asynchronous(true)]
- fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts(
- #[case] asynchronous: bool,
- ) {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let raised = PyValueError::new_err("over budget").into_value(py);
- let (result, log) = run_composed(
- py,
- success_machine(),
- SyntheticBinding {
- log: Log::default(),
- op: OpScript::Answer,
- classifier_fails: false,
- pending_reply: None,
- },
- HookScript::Plain,
- |hooks| {
- crate::HookChain::new()
- .with(hooks)
- .with(ArgumentPolicy::Reject(raised.clone_ref(py)))
- },
- call_options(asynchronous),
- );
- assert!(result.unwrap_err().value(py).is(raised.bind(py)));
- assert_eq!(
- log,
- [
- "started",
- "begin",
- "failed:Host:over budget",
- "adapter.close",
- "host.close"
- ]
- );
- });
- }
-
- #[rstest::rstest]
- #[case::synchronous(false)]
- #[case::asynchronous(true)]
- fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
- #[case] asynchronous: bool,
- ) {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- let (result, _) = run_composed(
- py,
- success_machine(),
- SyntheticBinding {
- log: Log::default(),
- op: OpScript::Answer,
- classifier_fails: false,
- pending_reply: None,
- },
- HookScript::Plain,
- |hooks| {
- crate::HookChain::new()
- .with(hooks)
- .with(ArgumentPolicy::Inherit)
- },
- call_options(asynchronous),
- );
- assert_eq!(
- result.unwrap().extract::(py).unwrap(),
- "project:2|sign|rewritten"
- );
- });
- }
-
- #[rstest::rstest]
- fn the_adapters_finalized_response_is_what_the_call_returns_and_reports() {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- for asynchronous in [false, true] {
- let (result, log) = run_scripted(
- py,
- success_machine(),
- OpScript::Answer,
- HookScript::ReplaceResponse,
- asynchronous,
- );
- assert_eq!(result.unwrap().extract::(py).unwrap(), "replaced");
- assert!(log.contains(&"succeeded:replaced".to_string()));
- assert!(!log.contains(&"succeeded:project:1|rewritten".to_string()));
- }
- });
- }
-
- #[rstest::rstest]
- fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding() {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- install_lifecycle_module(py);
- for asynchronous in [false, true] {
- let (result, log) = run_scripted(
- py,
- success_machine(),
- OpScript::Answer,
- HookScript::FailAfterSuccess,
- asynchronous,
- );
- let error = result.unwrap_err();
- assert_eq!(error.value(py).to_string(), "after_success failed");
- assert_eq!(
- &log[log.len() - 4..],
- [
- "after_success",
- "failed:Host:after_success failed",
- "adapter.close",
- "host.close"
- ]
- );
- assert!(!log.iter().any(|entry| entry.starts_with("succeeded")));
- }
- });
- }
-
- #[rstest::rstest]
- fn cancellation_ends_the_call_without_terminal_dispatch() {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- struct Cancelling(Log);
- impl PythonBinding for Cancelling {
- type Protocol = Synthetic;
- type Failure = Classified;
- fn decode_request(
- &mut self,
- _: Python<'_>,
- _: &Bound<'_, PyDict>,
- ) -> Result> {
- self.0.push("project");
- Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into())
- }
-
- fn encode_stream_head(
- &mut self,
- _: Python<'_>,
- head: std::convert::Infallible,
- ) -> PyResult> {
- match head {}
- }
- fn encode_chunk(
- &mut self,
- _: Python<'_>,
- chunk: std::convert::Infallible,
- ) -> PyResult> {
- match chunk {}
- }
- fn encode_response(&mut self, _: Python<'_>, _: String) -> PyResult> {
- Err(missing_state())
- }
- fn map_error(&self, _: Python<'_>, error: Error) -> PyResult {
- self.0.push("classify");
- Ok(Classified(error.0))
- }
- fn host_error(error: &PyErr) -> Error {
- Error(error.to_string())
- }
- }
-
- impl PythonHostCalls for Cancelling {
- fn handle_host_call(
- &mut self,
- _: Python<'_>,
- _: (&'static str, Reply),
- ) -> Result<(), InvokeError> {
- Err(missing_state().into())
- }
- }
-
- impl PythonOwned for Cancelling {
- fn close(&mut self, _: Python<'_>) {}
- fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
- Ok(())
- }
- }
- let log = Log::default();
- let host = Cancelling(Log(log.0.clone()));
- let adapter = SyntheticHooks {
- log: Log(log.0.clone()),
- script: HookScript::Plain,
- };
- let error = run_call(
- py,
- |_, _, request| Ok(success_machine()(request)),
- host,
- adapter,
- PyDict::new(py).unbind(),
- call_options(false),
- )
- .unwrap_err();
- assert!(!error.is_instance_of::(py));
- assert_eq!(
- log.entries(),
- ["started", "begin", "project", "adapter.close"]
- );
- });
- }
-
- #[rstest::rstest]
- fn python_driver_preserves_inline_await_and_native_ownership() {
- let _guard = PYTHON_GLOBALS
- .lock()
- .unwrap_or_else(|error| error.into_inner());
- crate::initialize_python();
- Python::attach(|py| {
- py.import("asyncio").unwrap();
- let module = install_lifecycle_module(py);
- let locals = PyDict::new(py);
- locals
- .set_item("drive", module.getattr("drive").unwrap())
- .unwrap();
- locals
- .set_item(
- "await_execution",
- wrap_pyfunction!(await_execution, py).unwrap(),
- )
- .unwrap();
- locals
- .set_item(
- "calling_execution",
- wrap_pyfunction!(calling_execution, py).unwrap(),
- )
- .unwrap();
- let probe = std::ffi::CString::new(include_str!("../tests/lifecycle.py")).unwrap();
- py.run(&probe, Some(&locals), Some(&locals)).unwrap();
- });
- }
- struct RetainingHost {
- retained: Option>,
- }
-
- impl ExecutionBody for RetainingHost {
- fn resume(&mut self, _: Option>>) -> PyResult {
- Python::attach(|py| Ok(ExecutionStep::Return(py.None())))
- }
-
- fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
- visit.call(&self.retained)
- }
- }
-
- #[pyfunction]
- fn retaining_coroutine(py: Python<'_>, retained: Py) -> PyResult> {
- Py::new(
- py,
- Execution::new(
- RetainingHost {
- retained: Some(retained),
- },
- lifecycle_binding,
- ),
- )
- }
-
- struct AwaitBody(Option>);
-
- impl ExecutionBody for AwaitBody {
- fn resume(&mut self, result: Option>>) -> PyResult {
- match self.0.take() {
- Some(awaitable) => Ok(ExecutionStep::Await(awaitable)),
- None => result
- .expect("selected await completed")
- .map(ExecutionStep::Return),
- }
- }
-
- fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
- visit.call(&self.0)
- }
- }
-
- #[pyfunction]
- fn await_execution(awaitable: Py) -> Execution {
- Execution::new(AwaitBody(Some(awaitable)), lifecycle_binding)
- }
-
- struct CallingBody(Py