From 76bf0cd5797a47b21c6d96ebd8ce645a197b519c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 23 Jul 2026 03:43:01 +0000 Subject: [PATCH 001/246] test(e2e): replace custom endpoints_client with provider SDK clients The llm_translation suite drove /v1/responses, /v1/messages, /embeddings, /v1/images/generations, /v1/moderations and /v1/audio/* through a bespoke endpoints_client wrapper that no customer runs. Tests now call the proxy the way customers do: the OpenAI SDK for the OpenAI-compatible surface and the Anthropic SDK for /v1/messages, wired through a session-scoped sdk fixture (sdk_clients.py) that points both SDKs at the proxy with a virtual key. Endpoints no official SDK covers keep the shared typed transport: rerank moves onto ProxyClient (RerankBody/RerankResponse in models.py) and the passthrough header test parses with the shared AnthropicMessagesResponse model. Files that only used endpoints_client for model registration now use the proxy fixture directly. endpoints_client.py is deleted; anthropic joins the e2e-dev dependency group so the lint env resolves the SDK imports. Resolves LIT-4577 --- pyproject.toml | 1 + tests/e2e/CLAUDE.md | 2 + tests/e2e/CONTRIBUTING.md | 2 + tests/e2e/llm_translation/conftest.py | 10 +- tests/e2e/llm_translation/endpoints_client.py | 385 ------------------ tests/e2e/llm_translation/sdk_clients.py | 53 +++ .../llm_translation/test_audio_speech_e2e.py | 99 ++--- .../test_audio_transcriptions_e2e.py | 25 +- .../e2e/llm_translation/test_cache_control.py | 5 +- .../test_credential_messages_e2e.py | 31 +- .../test_custom_pricing_e2e.py | 33 +- .../test_embeddings_endpoint_e2e.py | 92 ++--- .../test_image_generation_e2e.py | 73 ++-- .../test_messages_azure_foundry_e2e.py | 195 ++++----- .../e2e/llm_translation/test_messages_e2e.py | 176 ++++---- ...st_messages_mid_conversation_system_e2e.py | 124 +++--- ...onversation_system_native_providers_e2e.py | 136 ++++--- .../llm_translation/test_moderations_e2e.py | 61 +-- .../e2e/llm_translation/test_ocr_rust_e2e.py | 10 +- .../test_passthrough_headers_e2e.py | 8 +- tests/e2e/llm_translation/test_rerank_e2e.py | 48 +-- .../e2e/llm_translation/test_responses_e2e.py | 355 +++++++--------- .../test_responses_metadata_e2e.py | 69 ++-- tests/e2e/models.py | 19 + tests/e2e/proxy_client.py | 12 + uv.lock | 4 +- 26 files changed, 828 insertions(+), 1200 deletions(-) delete mode 100644 tests/e2e/llm_translation/endpoints_client.py create mode 100644 tests/e2e/llm_translation/sdk_clients.py diff --git a/pyproject.toml b/pyproject.toml index 62bd37c3db6..f152a6a73cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -198,6 +198,7 @@ e2e-dev = [ "playwright==1.61.0", "websockets>=15.0.1,<16.0", "locust==2.45.0", + "anthropic==0.84.0", ] proxy-dev = [ "prisma==0.11.0", diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0e39664e358..fa73cbacd96 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -58,6 +58,8 @@ That snippet only conveys intent. What you actually write uses the real harness: Every HTTP call goes through the shared transport, never through `requests.*` in a test. `e2e_http.py` is the only module permitted to call `requests.*`, and that is enforced in CI by `tests/code_coverage_tests/check_e2e_no_raw_requests.py`. A test that imports requests will fail the check +One deliberate exception: LLM-endpoint calls in `llm_translation/` go through the real provider SDKs (OpenAI, Anthropic) via the suite's `sdk` fixture (`llm_translation/sdk_clients.py`), because that is what customers actually run against the proxy (LIT-4577). The SDKs raise their own typed exceptions on failure, which is exactly the customer-observable contract; management routes (model/key CRUD, spend read-back) and endpoints no official SDK covers (e.g. `/v1/rerank`, `/v1/ocr`, custom passthrough paths) stay on the shared transport. Raw HTTP client imports remain banned either way + The shape is layered so tests stay declarative `transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index fc43769aca8..fa0174f686c 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -122,6 +122,8 @@ That snippet only conveys intent. What you actually write uses the real harness: Every HTTP call goes through the shared transport, never through `requests.*` in a test. `e2e_http.py` is the only module permitted to call `requests.*`, and that is enforced in CI by `tests/code_coverage_tests/check_e2e_no_raw_requests.py`. A test that imports requests will fail the check +One deliberate exception: LLM-endpoint calls in `llm_translation/` go through the real provider SDKs (OpenAI, Anthropic) via the suite's `sdk` fixture (`llm_translation/sdk_clients.py`), because that is what customers actually run against the proxy (LIT-4577). Management routes and endpoints no official SDK covers stay on the shared transport, and raw HTTP client imports remain banned either way + The shape is layered so tests stay declarative `transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test diff --git a/tests/e2e/llm_translation/conftest.py b/tests/e2e/llm_translation/conftest.py index f35ecf0760d..9fd45799773 100644 --- a/tests/e2e/llm_translation/conftest.py +++ b/tests/e2e/llm_translation/conftest.py @@ -2,14 +2,16 @@ The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. PassthroughClient holds the shared -ProxyClient, so the `resources` fixture cleans up keys this suite creates. +ProxyClient, so the `resources` fixture cleans up keys this suite creates. The +`sdk` fixture hands tests real provider SDK clients (OpenAI, Anthropic) pointed +at the proxy, the way customers actually call it. """ import pytest -from endpoints_client import EndpointsClient, build_endpoints_client from passthrough_client import PassthroughClient, build_client from proxy_client import ProxyClient +from sdk_clients import SdkClients, build_sdk_clients def pytest_configure(config: pytest.Config) -> None: @@ -25,5 +27,5 @@ def client(proxy: ProxyClient) -> PassthroughClient: @pytest.fixture(scope="session") -def endpoints_client(proxy: ProxyClient) -> EndpointsClient: - return build_endpoints_client(proxy) +def sdk() -> SdkClients: + return build_sdk_clients() diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py deleted file mode 100644 index ace621d03b3..00000000000 --- a/tests/e2e/llm_translation/endpoints_client.py +++ /dev/null @@ -1,385 +0,0 @@ -"""Client for the non-chat inference endpoints (responses, messages, rerank, -embeddings, audio speech, image generation). - -Each test registers the deployment it needs through /model/new (deleted on -teardown), so nothing is hardcoded into the gateway config, then drives the -endpoint with `send` and parses the provider-native body with a suite-local model -so the assertion is on real content, not just a 200. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Literal - -from pydantic import BaseModel - -from proxy_client import ProxyClient -from e2e_http import BinaryStream, Result, StreamingResponse -from models import CacheControl, ChatMessage, LiteLLMParamsBody, RichMessage, TextBlock - -__all__ = [ - "CacheControl", - "RichMessage", - "TextBlock", -] - - -class FunctionParameterProperty(BaseModel): - type: str - description: str | None = None - - -class FunctionParameters(BaseModel): - type: Literal["object"] = "object" - properties: dict[str, FunctionParameterProperty] - required: list[str] = [] - - -class ResponsesFunctionTool(BaseModel): - type: Literal["function"] = "function" - name: str - description: str | None = None - parameters: FunctionParameters - - -class ResponsesInputTextPart(BaseModel): - type: Literal["input_text"] = "input_text" - text: str - - -class ResponsesInputImagePart(BaseModel): - type: Literal["input_image"] = "input_image" - image_url: str - - -ResponsesInputContentPart = ResponsesInputTextPart | ResponsesInputImagePart - - -class ResponsesInputMessage(BaseModel): - role: Literal["user", "assistant", "system"] = "user" - content: list[ResponsesInputContentPart] - - -ResponsesInput = str | list[ResponsesInputMessage] - - -class ResponsesRequest(BaseModel): - model: str - input: ResponsesInput - instructions: str | None = None - stream: bool = False - tools: list[ResponsesFunctionTool] | None = None - - -class MessagesRequest(BaseModel): - model: str - max_tokens: int - messages: list[ChatMessage] - - -class RichMessagesRequest(BaseModel): - model: str - max_tokens: int = 64 - system: list[TextBlock] - messages: list[RichMessage] - - -class EmbeddingsRequest(BaseModel): - model: str - input: str - - -class RerankRequest(BaseModel): - model: str - query: str - documents: list[str] - top_n: int - - -class SpeechRequest(BaseModel): - model: str - input: str - voice: str - - -class ImageRequest(BaseModel): - model: str - prompt: str - n: int = 1 - size: str = "1024x1024" - - -class TranscriptionForm(BaseModel): - model: str - response_format: str = "json" - - -class ModerationRequest(BaseModel): - model: str - input: str - - -class ResponsesOutputContent(BaseModel): - type: str | None = None - text: str | None = None - - -class ResponsesOutputItem(BaseModel): - type: str | None = None - content: list[ResponsesOutputContent] = [] - name: str | None = None - arguments: str | None = None - call_id: str | None = None - - -class ResponsesResult(BaseModel): - id: str | None = None - status: str | None = None - model: str | None = None - output: list[ResponsesOutputItem] = [] - - @property - def text(self) -> str: - return "".join( - content.text or "" for item in self.output for content in item.content - ) - - @property - def function_calls(self) -> tuple[ResponsesOutputItem, ...]: - return tuple( - item - for item in self.output - if item.type == "function_call" - and item.name is not None - and item.arguments is not None - ) - - -class ResponsesStreamEvent(BaseModel): - event_id: str | None = None - - -class ResponsesStreamEventType(BaseModel): - type: str - - -class ResponsesOutputTextDeltaEvent(ResponsesStreamEvent): - type: Literal["response.output_text.delta"] - delta: str - - -class AnthropicContentBlock(BaseModel): - type: str | None = None - text: str | None = None - - -class MessagesUsage(BaseModel): - input_tokens: int = 0 - output_tokens: int = 0 - cache_creation_input_tokens: int = 0 - cache_read_input_tokens: int = 0 - - -class MessagesResult(BaseModel): - id: str | None = None - role: str | None = None - model: str | None = None - content: list[AnthropicContentBlock] = [] - usage: MessagesUsage = MessagesUsage() - - @property - def text(self) -> str: - return "".join(block.text or "" for block in self.content) - - -class EmbeddingItem(BaseModel): - embedding: list[float] = [] - - -class EmbeddingsResult(BaseModel): - data: list[EmbeddingItem] = [] - - @property - def first_vector(self) -> tuple[float, ...]: - return tuple(self.data[0].embedding) if self.data else () - - -class RerankItem(BaseModel): - index: int | None = None - relevance_score: float | None = None - - -class RerankResult(BaseModel): - results: list[RerankItem] = [] - - -class ImageItem(BaseModel): - url: str | None = None - b64_json: str | None = None - - -class ImagesResult(BaseModel): - data: list[ImageItem] = [] - - -class TranscriptionResult(BaseModel): - text: str = "" - - -class ModerationResultItem(BaseModel): - flagged: bool - categories: dict[str, bool] = {} - - @property - def flagged_categories(self) -> tuple[str, ...]: - return tuple(name for name, hit in self.categories.items() if hit) - - -class ModerationResult(BaseModel): - results: list[ModerationResultItem] = [] - - @property - def first(self) -> ModerationResultItem | None: - return self.results[0] if self.results else None - - -@dataclass(frozen=True, slots=True) -class EndpointsClient: - proxy: ProxyClient - - def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str: - return self.proxy.create_model(model_name, litellm_params) - - def delete_model(self, model_id: str) -> None: - self.proxy.delete_model(model_id) - - def _send( - self, path: str, key: str, body: BaseModel, *, stream: bool = False - ) -> StreamingResponse: - return self.proxy.transport.send( - path, - headers=self.proxy.transport.bearer(key), - json=body, - stream=stream, - ) - - def responses( - self, key: str, model: str, text: str, *, stream: bool = False - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=text, - instructions="You are a helpful assistant", - stream=stream, - ), - stream=stream, - ) - - def responses_vision( - self, key: str, model: str, text: str, image_url: str - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=[ - ResponsesInputMessage( - content=[ - ResponsesInputTextPart(text=text), - ResponsesInputImagePart(image_url=image_url), - ] - ) - ], - instructions="You are a helpful assistant", - ), - ) - - def responses_with_tools( - self, key: str, model: str, text: str, tools: list[ResponsesFunctionTool] - ) -> StreamingResponse: - return self._send( - "/v1/responses", - key, - ResponsesRequest( - model=model, - input=text, - instructions="You are a helpful assistant", - tools=tools, - ), - ) - - def messages( - self, key: str, model: str, text: str, *, max_tokens: int = 64 - ) -> StreamingResponse: - return self._send( - "/v1/messages", - key, - MessagesRequest( - model=model, - max_tokens=max_tokens, - messages=[ChatMessage(role="user", content=text)], - ), - ) - - def embeddings(self, key: str, model: str, text: str) -> StreamingResponse: - return self._send("/embeddings", key, EmbeddingsRequest(model=model, input=text)) - - def rerank( - self, key: str, model: str, query: str, documents: list[str], top_n: int - ) -> StreamingResponse: - return self._send( - "/v1/rerank", - key, - RerankRequest(model=model, query=query, documents=documents, top_n=top_n), - ) - - def audio_speech( - self, key: str, model: str, text: str, *, voice: str = "alloy" - ) -> StreamingResponse: - return self._send( - "/v1/audio/speech", key, SpeechRequest(model=model, input=text, voice=voice) - ) - - def audio_speech_stream( - self, key: str, model: str, text: str, *, voice: str = "alloy" - ) -> BinaryStream: - return self.proxy.transport.stream_binary( - "/v1/audio/speech", - headers=self.proxy.transport.bearer(key), - json=SpeechRequest(model=model, input=text, voice=voice), - ) - - def transcribe( - self, key: str, model: str, *, filename: str, content: bytes - ) -> Result[TranscriptionResult]: - return self.proxy.transport.upload( - "/v1/audio/transcriptions", - headers=self.proxy.transport.bearer(key), - form=TranscriptionForm(model=model), - filename=filename, - content=content, - file_content_type="audio/wav", - response_type=TranscriptionResult, - ) - - def moderations(self, key: str, model: str, text: str) -> Result[ModerationResult]: - return self.proxy.transport.post( - "/v1/moderations", - headers=self.proxy.transport.bearer(key), - json=ModerationRequest(model=model, input=text), - response_type=ModerationResult, - ) - - def images(self, key: str, model: str, prompt: str) -> StreamingResponse: - return self._send( - "/v1/images/generations", key, ImageRequest(model=model, prompt=prompt) - ) - - -def build_endpoints_client(proxy: ProxyClient) -> EndpointsClient: - return EndpointsClient(proxy=proxy) diff --git a/tests/e2e/llm_translation/sdk_clients.py b/tests/e2e/llm_translation/sdk_clients.py new file mode 100644 index 00000000000..8a6b187f979 --- /dev/null +++ b/tests/e2e/llm_translation/sdk_clients.py @@ -0,0 +1,53 @@ +"""Real provider SDK clients pointed at the proxy, connected the way customers +connect (LIT-4577). + +The OpenAI SDK drives the OpenAI-compatible surface (/responses, /embeddings, +/images/generations, /moderations, /audio/*) and the Anthropic SDK drives +/v1/messages, each authenticated with a litellm virtual key. Errors surface as +the SDK's own exceptions, exactly what an end user sees. Retries are disabled +so a proxy fault fails the test instead of being papered over, and the timeout +matches the shared transport's request budget. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +from anthropic import Anthropic +from openai import OpenAI + +from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT + + +def response_header(headers: Mapping[str, str], name: str) -> str | None: + """Typed read of an SDK response header: httpx.Headers.get returns Any and + httpx itself is a banned import in suite code, so tests read headers through + the Mapping[str, str] interface Headers fulfils.""" + return headers[name] if name in headers else None + + +@dataclass(frozen=True, slots=True) +class SdkClients: + base_url: str + request_timeout: float + + def openai(self, key: str) -> OpenAI: + return OpenAI( + base_url=self.base_url, + api_key=key, + timeout=self.request_timeout, + max_retries=0, + ) + + def anthropic(self, key: str) -> Anthropic: + return Anthropic( + base_url=self.base_url, + api_key=key, + timeout=self.request_timeout, + max_retries=0, + ) + + +def build_sdk_clients() -> SdkClients: + return SdkClients(base_url=PROXY_BASE_URL, request_timeout=REQUEST_TIMEOUT) diff --git a/tests/e2e/llm_translation/test_audio_speech_e2e.py b/tests/e2e/llm_translation/test_audio_speech_e2e.py index b95cef8db4d..859292ecc7b 100644 --- a/tests/e2e/llm_translation/test_audio_speech_e2e.py +++ b/tests/e2e/llm_translation/test_audio_speech_e2e.py @@ -1,9 +1,10 @@ """Live e2e: POST /v1/audio/speech returns audio, non-streamed and streamed. -The non-streamed call asserts an audio (not JSON) body. The streamed call consumes -the response the way a player would and asserts customer-observable streaming: -chunked transfer encoding (a buffered body would carry a content-length) with -non-zero audio bytes. +Both calls go through the real OpenAI SDK (LIT-4577). The non-streamed call +asserts an audio (not JSON) body. The streamed call consumes the response the +way a player would and asserts customer-observable streaming: chunked transfer +encoding (a buffered body would carry a content-length) with non-zero audio +bytes. """ from __future__ import annotations @@ -11,68 +12,70 @@ from __future__ import annotations import pytest from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients, response_header pytestmark = pytest.mark.e2e +def _register(proxy: ProxyClient, resources: ResourceManager, prefix: str) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model( + model, + LiteLLMParamsBody(model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY"), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + return model + + class TestAudioSpeech: @pytest.mark.covers("llm.audio_speech.openai.basic.nonstream.works") def test_audio_speech_returns_audio( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-speech-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, "e2e-speech") + client = sdk.openai(resources.key()) - result = endpoints_client.audio_speech(key, model, "Hello!") - require_successful_call(result) - assert "audio" in (result.content_type or ""), ( - f"/audio/speech content-type is not audio: {result.content_type!r}" + response = client.audio.speech.with_raw_response.create( + model=model, voice="alloy", input="Hello!" ) - assert result.body, "/audio/speech returned an empty body" + content_type = response_header(response.headers, "content-type") + assert "audio" in (content_type or ""), ( + f"/audio/speech content-type is not audio: {content_type!r}" + ) + assert response.content, "/audio/speech returned an empty body" @pytest.mark.covers("llm.audio_speech.openai.basic.stream.works") def test_audio_speech_streams_audio_chunks( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-speech-stream-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="openai/gpt-4o-mini-tts", api_key="os.environ/OPENAI_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, "e2e-speech-stream") + client = sdk.openai(resources.key()) - result = endpoints_client.audio_speech_stream( - key, - model, - "Streaming speech should arrive in several audio chunks so a client can " - "begin playback well before the whole clip has finished generating.", + with client.audio.speech.with_streaming_response.create( + model=model, + voice="alloy", + input=( + "Streaming speech should arrive in several audio chunks so a client can " + "begin playback well before the whole clip has finished generating." + ), + ) as response: + content_type = response_header(response.headers, "content-type") + transfer_encoding = response_header(response.headers, "transfer-encoding") + content_length = response_header(response.headers, "content-length") + total_bytes = sum(len(chunk) for chunk in response.iter_bytes(chunk_size=8192)) + + assert "audio" in (content_type or ""), ( + f"/audio/speech content-type is not audio: {content_type!r}" ) - assert result.ok, ( - f"/audio/speech stream failed (status {result.status_code}); body={result.error_body}" + assert "chunked" in (transfer_encoding or ""), ( + f"/audio/speech did not stream: transfer-encoding={transfer_encoding!r}, " + f"content-length={content_length!r} (a buffered body is not a stream)" ) - assert "audio" in (result.content_type or ""), ( - f"/audio/speech content-type is not audio: {result.content_type!r}" - ) - assert result.chunked, ( - f"/audio/speech did not stream: transfer-encoding={result.transfer_encoding!r}, " - f"content-length={result.content_length!r} (a buffered body is not a stream)" - ) - assert result.content_length is None, ( - f"/audio/speech advertised content-length={result.content_length!r} on a " + assert content_length is None, ( + f"/audio/speech advertised content-length={content_length!r} on a " f"streamed response (a buffered body is not a stream)" ) - assert result.total_bytes > 0, "/audio/speech stream returned no audio bytes" + assert total_bytes > 0, "/audio/speech stream returned no audio bytes" diff --git a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py index af6123dc46a..617192f88c7 100644 --- a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py +++ b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py @@ -1,8 +1,9 @@ """Live e2e: POST /v1/audio/transcriptions turns speech into text. Registers an OpenAI speech-to-text deployment at runtime and uploads a spoken -weather question (the realtime suite's 24kHz WAV fixture) as multipart, asserting -the returned transcript is non-empty and mentions the word it was asked about. +weather question (the realtime suite's 24kHz WAV fixture) through the real +OpenAI SDK (LIT-4577), asserting the returned transcript is non-empty and +mentions the word it was asked about. """ from __future__ import annotations @@ -12,10 +13,10 @@ from pathlib import Path import pytest from e2e_config import unique_marker -from e2e_http import unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -27,24 +28,22 @@ WEATHER_WAV = ( class TestAudioTranscriptions: @pytest.mark.covers("llm.audio_transcriptions.openai.basic.nonstream.works") def test_audio_transcriptions_returns_text( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: model = f"e2e-transcribe-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="openai/gpt-4o-mini-transcribe", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) - result = unwrap( - endpoints_client.transcribe( - key, model, filename=WEATHER_WAV.name, content=WEATHER_WAV.read_bytes() - ) + transcription = client.audio.transcriptions.create( + model=model, file=(WEATHER_WAV.name, WEATHER_WAV.read_bytes(), "audio/wav") ) - text = result.text.strip() + text = transcription.text.strip() assert text, "/audio/transcriptions returned an empty transcript" assert "weather" in text.lower(), ( f"transcript of a spoken weather question does not mention weather: {text!r}" diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index a2c17b0fb66..888cbb422eb 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -15,7 +15,7 @@ service_tier lives in test_provider_features_e2e.py. The provider-native cache_control request shape is not expressible with the shared ``ChatBody`` (whose content is a plain string), so the cacheable body is -built from the typed content blocks shared in ``endpoints_client.py``. +built from the typed content blocks shared in ``models.py``. """ from __future__ import annotations @@ -27,9 +27,8 @@ from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import Result, unwrap -from endpoints_client import CacheControl, RichMessage, TextBlock from lifecycle import ResourceManager -from models import ChatResponse, LiteLLMParamsBody, Usage +from models import CacheControl, ChatResponse, LiteLLMParamsBody, RichMessage, TextBlock, Usage from passthrough_client import PassthroughClient import os diff --git a/tests/e2e/llm_translation/test_credential_messages_e2e.py b/tests/e2e/llm_translation/test_credential_messages_e2e.py index 49ea748430e..d4bf4566cc0 100644 --- a/tests/e2e/llm_translation/test_credential_messages_e2e.py +++ b/tests/e2e/llm_translation/test_credential_messages_e2e.py @@ -7,43 +7,48 @@ import os import pytest from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager from models import CredentialCreateBody, LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e class TestCredentialBackedMessages: @pytest.mark.covers("mgmt.credential.new.serves_request") - def test_credential_backed_messages(self, endpoints_client: EndpointsClient, resources: ResourceManager) -> None: + def test_credential_backed_messages( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: marker = unique_marker() credential_name = f"e2e-cred-{marker}" model = f"e2e-cred-messages-{marker}" anthropic_api_key = os.getenv("ANTHROPIC_API_KEY") assert anthropic_api_key, "ANTHROPIC_API_KEY must be set for this live e2e test" - endpoints_client.proxy.create_credential( + proxy.create_credential( CredentialCreateBody( credential_name=credential_name, credential_values={"api_key": anthropic_api_key}, ) ) - resources.defer(lambda: endpoints_client.proxy.delete_credential(credential_name)) + resources.defer(lambda: proxy.delete_credential(credential_name)) - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="anthropic/claude-haiku-4-5", litellm_credential_name=credential_name, ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) - key = resources.key() - result = endpoints_client.messages(key, model, "reply with one word") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}" - assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}" + client = sdk.anthropic(resources.key()) + message = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "reply with one word"}], + ) + assert message.role == "assistant", f"unexpected role: {message.role!r}" + text = "".join(block.text for block in message.content if block.type == "text") + assert text.strip(), f"/v1/messages returned no text: {message.content!r}" diff --git a/tests/e2e/llm_translation/test_custom_pricing_e2e.py b/tests/e2e/llm_translation/test_custom_pricing_e2e.py index b4ff631a56b..1cebf90fa21 100644 --- a/tests/e2e/llm_translation/test_custom_pricing_e2e.py +++ b/tests/e2e/llm_translation/test_custom_pricing_e2e.py @@ -25,7 +25,6 @@ from pydantic import BaseModel, RootModel from e2e_config import unique_marker from proxy_client import ProxyClient from e2e_http import Success, unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import ( ChatBody, @@ -71,7 +70,7 @@ def _approx_equal(actual: float, expected: float) -> bool: def _provision( - endpoints_client: EndpointsClient, + proxy: ProxyClient, resources: ResourceManager, prefix: str, *, @@ -84,7 +83,7 @@ def _provision( marker keeps the name unique so concurrent runs on the shared proxy never collide.""" model_name = f"{prefix}-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model_name, LiteLLMParamsBody( model=BACKEND_MODEL, @@ -93,15 +92,15 @@ def _provision( output_cost_per_token=output_cost_per_token, ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model_name def _provision_custom_priced( - endpoints_client: EndpointsClient, resources: ResourceManager + proxy: ProxyClient, resources: ResourceManager ) -> str: return _provision( - endpoints_client, + proxy, resources, "custom-priced-flash", input_cost_per_token=CUSTOM_INPUT_RATE, @@ -151,14 +150,14 @@ def _poll_breakdown_row(proxy: ProxyClient, key: str, response_id: str | None) - class TestCustomPricing: def test_custom_pricing_is_billed_at_configured_rate( self, - endpoints_client: EndpointsClient, + proxy: ProxyClient, resources: ResourceManager, scoped_key: str, ) -> None: - model = _provision_custom_priced(endpoints_client, resources) + model = _provision_custom_priced(proxy, resources) chat = unwrap( - endpoints_client.proxy.chat( + proxy.chat( scoped_key, ChatBody( model=model, @@ -172,7 +171,7 @@ class TestCustomPricing: ) ) - row = _poll_breakdown_row(endpoints_client.proxy, scoped_key, chat.id) + row = _poll_breakdown_row(proxy, scoped_key, chat.id) assert row.metadata and row.metadata.cost_breakdown # guaranteed by the poll breakdown = row.metadata.cost_breakdown @@ -195,10 +194,10 @@ class TestCustomPricing: ) def test_model_info_reports_custom_pricing( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: - model = _provision_custom_priced(endpoints_client, resources) - entry = _model_info_entry(endpoints_client.proxy.model_info(), model) + model = _provision_custom_priced(proxy, resources) + entry = _model_info_entry(proxy.model_info(), model) assert entry.litellm_params.input_cost_per_token == CUSTOM_INPUT_RATE, ( f"/model/info litellm_params input rate " @@ -210,20 +209,20 @@ class TestCustomPricing: ) def test_custom_pricing_is_isolated_from_sibling_deployment( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: # Register the override first so its rate is in the backend cost map before # the sibling resolves; a leak (LIT-3897) would then poison the sibling. - custom = _provision_custom_priced(endpoints_client, resources) + custom = _provision_custom_priced(proxy, resources) sibling = _provision( - endpoints_client, + proxy, resources, "base-flash", input_cost_per_token=None, output_cost_per_token=None, ) - entries = {entry.model_name: entry for entry in endpoints_client.proxy.model_info()} + entries = {entry.model_name: entry for entry in proxy.model_info()} custom_entry = entries.get(custom) sibling_entry = entries.get(sibling) assert custom_entry is not None, f"{custom} absent from /model/info" diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index 157caedd561..dac69591341 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -1,9 +1,10 @@ """Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex. -Each test registers the deployment it needs at runtime (deleted on teardown) and -asserts a non-empty, non-zero vector came back. The LIT-3167 guard in -tests/e2e/embeddings/ covers the Gemini embedding path; embeddings cost tracking is -covered by tests/e2e/quota_management/spend_tracking/. +Each test registers the deployment it needs at runtime (deleted on teardown), +drives the endpoint with the real OpenAI SDK (LIT-4577), and asserts a +non-empty, non-zero vector came back. The LIT-3167 guard in +tests/e2e/embeddings/ covers the Gemini embedding path; embeddings cost tracking +is covered by tests/e2e/quota_management/spend_tracking/. """ from __future__ import annotations @@ -11,79 +12,74 @@ from __future__ import annotations import pytest from e2e_config import unique_marker -from e2e_http import require_successful_call -from endpoints_client import EmbeddingsResult, EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e +def _assert_embedding_vector( + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + prefix: str, + params: LiteLLMParamsBody, +) -> None: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) + + embeddings = client.embeddings.create(model=model, input="Say this is a test!") + assert embeddings.data, f"/embeddings returned no data: {embeddings!r}" + vector = embeddings.data[0].embedding + assert vector, f"/embeddings returned no vector: {embeddings!r}" + assert any(component != 0.0 for component in vector), "embedding vector is all zeros" + + class TestEmbeddingsEndpoint: @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") def test_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings", LiteLLMParamsBody( model="openai/text-embedding-3-small", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) @pytest.mark.covers("llm.embeddings.bedrock.basic.nonstream.works") def test_bedrock_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-bedrock-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings-bedrock", LiteLLMParamsBody( model="bedrock/amazon.titan-embed-text-v2:0", aws_region_name="us-west-2" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) @pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works") def test_vertex_embeddings_returns_vector( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-embeddings-vertex-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_embedding_vector( + proxy, + resources, + sdk, + "e2e-embeddings-vertex", LiteLLMParamsBody( model="vertex_ai/gemini-embedding-2", vertex_project="os.environ/VERTEXAI_PROJECT", vertex_location="us-central1", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.embeddings(key, model, "Say this is a test!") - require_successful_call(result) - parsed = EmbeddingsResult.model_validate_json(result.body) - assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" - assert any(component != 0.0 for component in parsed.first_vector), ( - f"embedding vector is all zeros: {result.body[:300]}" - ) diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index 1ba78a7e083..c5204162501 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -1,8 +1,7 @@ """Live e2e: POST /v1/images/generations returns an image. -Registers an OpenAI image deployment at runtime and asserts the response carries a -generated image (url or base64). Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +Registers an image deployment at runtime, drives it through the real OpenAI SDK +(LIT-4577), and asserts the response carries a generated image (url or base64). """ from __future__ import annotations @@ -10,50 +9,58 @@ from __future__ import annotations import pytest from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, ImagesResult from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e -def _assert_image_returned(body: str) -> None: - parsed = ImagesResult.model_validate_json(body) - assert parsed.data, f"/images/generations returned no data: {body[:300]}" - first = parsed.data[0] - assert first.b64_json or first.url, ( - f"generated image has neither b64_json nor url: {body[:300]}" - ) +def _assert_image_returned( + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + prefix: str, + params: LiteLLMParamsBody, +) -> None: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) + + images = client.images.generate(model=model, prompt="Draw a cute cat", n=1, size="1024x1024") + data = images.data or [] + assert data, f"/images/generations returned no data: {images!r}" + first = data[0] + assert first.b64_json or first.url, "generated image has neither b64_json nor url" class TestImageGeneration: @pytest.mark.covers("llm.images_generations.openai.basic.nonstream.works") def test_image_generation_returns_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-image-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="openai/gpt-image-1-mini", api_key="os.environ/OPENAI_API_KEY" - ), + _assert_image_returned( + proxy, + resources, + sdk, + "e2e-image", + LiteLLMParamsBody(model="openai/gpt-image-1-mini", api_key="os.environ/OPENAI_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - result = endpoints_client.images(key, model, "Draw a cute cat") - require_successful_call(result) - _assert_image_returned(result.body) - - @pytest.mark.covers("llm.images_generations.bedrock.basic.nonstream.works", exercised_on=["images_generations"]) + @pytest.mark.covers( + "llm.images_generations.bedrock.basic.nonstream.works", exercised_on=["images_generations"] + ) def test_bedrock_image_generation_returns_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") - model = f"e2e-bedrock-image-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + _assert_image_returned( + proxy, + resources, + sdk, + "e2e-bedrock-image", LiteLLMParamsBody( model="bedrock/amazon.titan-image-generator-v2:0", aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", @@ -61,9 +68,3 @@ class TestImageGeneration: aws_region_name="os.environ/AWS_REGION", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - - result = endpoints_client.images(key, model, "Draw a cute cat") - require_successful_call(result) - _assert_image_returned(result.body) diff --git a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py index d8d44820e80..90568d8c951 100644 --- a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py +++ b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py @@ -1,9 +1,9 @@ """Live e2e: POST /v1/messages routed to Azure AI Foundry Anthropic deployments. Registers `azure_ai/` deployments at runtime and drives the Messages -endpoint through the gateway across the behaviors an Anthropic client relies on: -a basic completion, a streamed completion, and tool use (non-streaming and -streaming). Auth is the Azure API key (`x-api-key`); the deployment reads +endpoint through the gateway with the real Anthropic SDK (LIT-4577) across the +behaviors an Anthropic client relies on: a basic completion, a streamed +completion, and tool use (non-streaming and streaming). The deployment reads `AZURE_AI_API_BASE` / `AZURE_AI_API_KEY` from the proxy env, so no secret is sent in the request. """ @@ -11,60 +11,50 @@ sent in the request. from __future__ import annotations import pytest +from anthropic.types import RawMessageStreamEvent, ToolParam from e2e_config import EXPECT_RUST, unique_marker -from e2e_http import StreamingResponse, require_successful_call, unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager -from models import ( - AnthropicCustomTool, - AnthropicMessagesBody, - ChatMessage, - JsonSchemaProperty, - LiteLLMParamsBody, - ToolInputSchema, -) +from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e AZURE_FOUNDRY_MODEL = "azure_ai/claude-haiku-4-5" -WEATHER_TOOL = AnthropicCustomTool( - name="get_weather", - description="Get the current weather for a city.", - input_schema=ToolInputSchema( - properties={"city": JsonSchemaProperty(type="string")}, - required=["city"], - ), -) +WEATHER_TOOL: ToolParam = { + "name": "get_weather", + "description": "Get the current weather for a city.", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, +} -def _assert_streamed_ok(result: StreamingResponse) -> None: - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" - assert any("content_block_delta" in event for event in result.stream_events), ( - "stream carried no content deltas" +def _assert_rust_served(headers: dict[str, str]) -> None: + if not EXPECT_RUST: + return + assert headers.get("x-litellm-rust") == "true", ( + "E2E_EXPECT_RUST is set, so this gateway must serve /v1/messages through the " + "Rust path, but the response carried no x-litellm-rust marker. The request " + "still succeeded, which is exactly the failure mode: a gateway whose native " + f"extension is unavailable falls back to Python silently. headers={headers}" ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" - ) - if EXPECT_RUST: - assert result.headers.get("x-litellm-rust") == "true", ( - "E2E_EXPECT_RUST is set, so this gateway must serve /v1/messages through the " - "Rust path, but the response carried no x-litellm-rust marker. The request " - "still succeeded, which is exactly the failure mode: a gateway whose native " - f"extension is unavailable falls back to Python silently. headers={result.headers}" - ) + + +def _assert_streamed_ok(event_types: list[str]) -> None: + assert event_types, "stream produced no SSE events" + assert "content_block_delta" in event_types, "stream carried no content deltas" + assert "message_stop" in event_types, "stream never reached message_stop" class TestAzureFoundryMessages: - def _register( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> tuple[str, str]: + def _register(self, proxy: ProxyClient, resources: ResourceManager) -> str: model = f"e2e-azure-foundry-messages-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model=AZURE_FOUNDRY_MODEL, @@ -72,91 +62,78 @@ class TestAzureFoundryMessages: api_key="os.environ/AZURE_AI_API_KEY", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - return model, resources.key(models=[model]) + resources.defer(lambda: proxy.delete_model(model_id)) + return model @pytest.mark.covers("llm.messages.azure_foundry.basic.nonstream.works") def test_basic_nonstream( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=64, - messages=[ChatMessage(role="user", content="Reply with one word.")], - ), - ) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + message = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "Reply with one word."}], ) - assert response.content, f"no content blocks in response: {response}" - text = "".join(block.text or "" for block in response.content if block.type == "text") - assert text.strip(), f"/v1/messages returned no text: {response}" + assert message.content, f"no content blocks in response: {message!r}" + text = "".join(block.text for block in message.content if block.type == "text") + assert text.strip(), f"/v1/messages returned no text: {message.content!r}" @pytest.mark.covers("llm.messages.azure_foundry.basic.stream.works") def test_basic_stream( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=64, - stream=True, - messages=[ChatMessage(role="user", content="Count from one to three.")], - ), + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + raw = client.messages.with_raw_response.create( + model=model, + max_tokens=64, + stream=True, + messages=[{"role": "user", "content": "Count from one to three."}], ) - _assert_streamed_ok(result) + _assert_rust_served({name.lower(): value for name, value in raw.headers.items()}) + _assert_streamed_ok([event.type for event in raw.parse()]) @pytest.mark.covers("llm.messages.azure_foundry.tool_use.nonstream.works") def test_tool_use_nonstream( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + message = client.messages.create( + model=model, + max_tokens=256, + tools=[WEATHER_TOOL], + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], ) - assert response.content, f"no content blocks in response: {response}" - assert any(block.type == "tool_use" for block in response.content), ( - f"model did not call the tool: {response}" + assert message.content, f"no content blocks in response: {message!r}" + assert any(block.type == "tool_use" for block in message.content), ( + f"model did not call the tool: {message.content!r}" ) @pytest.mark.covers("llm.messages.azure_foundry.tool_use.stream.works") def test_tool_use_stream( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - stream=True, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" - assert any("tool_use" in event for event in result.stream_events), ( - "stream carried no tool_use block" - ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key(models=[model])) + + stream = client.messages.create( + model=model, + max_tokens=256, + stream=True, + tools=[WEATHER_TOOL], + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], ) + events: list[RawMessageStreamEvent] = list(stream) + event_types = [event.type for event in events] + assert event_types, "stream produced no SSE events" + assert any( + event.type == "content_block_start" and event.content_block.type == "tool_use" + for event in events + ), "stream carried no tool_use block" + assert "message_stop" in event_types, "stream never reached message_stop" diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 44376218c6b..0c5734bc299 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -1,41 +1,35 @@ """Live e2e: POST /v1/messages (Anthropic Messages API) returns a real completion. -Registers an Anthropic deployment at runtime, drives the Messages endpoint through -the gateway, and asserts an assistant message with text came back, both -non-streaming and streamed. Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +Registers an Anthropic deployment at runtime and drives the Messages endpoint +through the gateway with the real Anthropic SDK, the client customers actually +use (LIT-4577), asserting an assistant message with text came back, both +non-streaming and streamed. """ from __future__ import annotations import pytest +from anthropic.types import Message, ToolParam from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call, unwrap -from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager -from models import ( - AnthropicCustomTool, - AnthropicMessagesBody, - ChatMessage, - JsonSchemaProperty, - LiteLLMParamsBody, - SpendLogRow, - ToolInputSchema, -) +from models import LiteLLMParamsBody, SpendLogRow +from proxy_client import ProxyClient +from sdk_clients import SdkClients, response_header pytestmark = pytest.mark.e2e ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5" -WEATHER_TOOL = AnthropicCustomTool( - name="get_weather", - description="Get the current weather for a city.", - input_schema=ToolInputSchema( - properties={"city": JsonSchemaProperty(type="string")}, - required=["city"], - ), -) +WEATHER_TOOL: ToolParam = { + "name": "get_weather", + "description": "Get the current weather for a city.", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, +} def _approx_equal(actual: float, expected: float) -> bool: @@ -43,60 +37,66 @@ def _approx_equal(actual: float, expected: float) -> bool: return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if block.type == "text") + + class TestAnthropicMessages: def _register( - self, endpoints_client: EndpointsClient, resources: ResourceManager - ) -> tuple[str, str]: - model = f"e2e-messages-{unique_marker()}" - model_id = endpoints_client.create_model( + self, proxy: ProxyClient, resources: ResourceManager, prefix: str = "e2e-messages" + ) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = proxy.create_model( model, - LiteLLMParamsBody( - model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" - ), + LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - return model, resources.key() + resources.defer(lambda: proxy.delete_model(model_id)) + return model @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.works") def test_messages_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key()) - result = endpoints_client.messages(key, model, "reply with one word") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}" - assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}" + message = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "reply with one word"}], + ) + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert _text(message).strip(), f"/v1/messages returned no text: {message.content!r}" @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.cost_logged") def test_messages_logs_cost_matching_the_response_header( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: require_env("ANTHROPIC_API_KEY") - model = f"e2e-messages-cost-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + model = self._register(proxy, resources, prefix="e2e-messages-cost") key = resources.key() + client = sdk.anthropic(key) - result = endpoints_client.messages(key, model, f"reply with one word {unique_marker()}") - require_successful_call(result) - parsed = MessagesResult.model_validate_json(result.body) - assert parsed.role == "assistant" and parsed.text.strip(), ( - f"/v1/messages returned no assistant text: {result.body[:300]}" + raw = client.messages.with_raw_response.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": f"reply with one word {unique_marker()}"}], + ) + message = raw.parse() + assert message.role == "assistant" and _text(message).strip(), ( + f"/v1/messages returned no assistant text: {message.content!r}" ) # The customer reads per-request cost off the response header (LIT-4076), so # it must be present and positive on /v1/messages, not only /chat/completions. - header_cost = result.response_cost - assert header_cost is not None and header_cost > 0, ( - "x-litellm-response-cost header missing or non-positive on /v1/messages; " - f"headers={result.headers}" + raw_header_cost = response_header(raw.headers, "x-litellm-response-cost") + assert raw_header_cost is not None, ( + "x-litellm-response-cost header missing on /v1/messages; " + f"headers={dict(raw.headers)}" + ) + header_cost = float(raw_header_cost) + assert header_cost > 0, ( + f"x-litellm-response-cost header non-positive on /v1/messages: {header_cost}" ) # Correlate the spend row by the unique scoped key, not the Anthropic response @@ -107,7 +107,7 @@ class TestAnthropicMessages: def _priced(rows: list[SpendLogRow]) -> bool: return any(r.spend is not None and r.spend > 0 for r in rows) - rows = endpoints_client.proxy.poll_logs_for_key(key, predicate=_priced) + rows = proxy.poll_logs_for_key(key, predicate=_priced) priced = [r for r in rows if r.spend is not None and r.spend > 0] assert priced, ( f"no priced /spend/logs row landed for key {key} within the poll window; got {rows}" @@ -123,50 +123,36 @@ class TestAnthropicMessages: @pytest.mark.covers("llm.messages.anthropic.basic.stream.works") def test_messages_streams_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key()) - result = endpoints_client.proxy.messages_stream( - key, - AnthropicMessagesBody( - model=model, - max_tokens=64, - stream=True, - messages=[ChatMessage(role="user", content="Count from one to three.")], - ), - ) - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_events, "stream produced no SSE events" - assert any("content_block_delta" in event for event in result.stream_events), ( - "stream carried no content deltas" - ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" + stream = client.messages.create( + model=model, + max_tokens=64, + stream=True, + messages=[{"role": "user", "content": "Count from one to three."}], ) + event_types = [event.type for event in stream] + assert event_types, "stream produced no SSE events" + assert "content_block_delta" in event_types, "stream carried no content deltas" + assert "message_stop" in event_types, "stream never reached message_stop" @pytest.mark.covers("llm.messages.anthropic.tool_use.nonstream.works") def test_messages_tool_use( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model, key = self._register(endpoints_client, resources) + model = self._register(proxy, resources) + client = sdk.anthropic(resources.key()) - response = unwrap( - endpoints_client.proxy.messages( - key, - AnthropicMessagesBody( - model=model, - max_tokens=256, - tools=[WEATHER_TOOL], - messages=[ - ChatMessage(role="user", content="What is the weather in Paris? Use the tool.") - ], - ), - ) + message = client.messages.create( + model=model, + max_tokens=256, + tools=[WEATHER_TOOL], + messages=[{"role": "user", "content": "What is the weather in Paris? Use the tool."}], ) - assert response.content, f"no content blocks in response: {response}" - assert any(block.type == "tool_use" for block in response.content), ( - f"model did not call the tool: {response}" + assert message.content, f"no content blocks in response: {message!r}" + assert any(block.type == "tool_use" for block in message.content), ( + f"model did not call the tool: {message.content!r}" ) diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index 4b3191e60bb..b79f3027692 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -17,27 +17,28 @@ entry whose prefix spans ``system`` plus message turns is invalidated when the reminder is hoisted (the ``system`` field mutates and a turn disappears from ``messages``), while an entry ending at the system block itself would survive the hoist and mask the regression. + +Calls go through the real Anthropic SDK (LIT-4577). The SDK's ``MessageParam`` +type only admits user/assistant roles, so the system reminder turn is cast to +it; the SDK serializes the dict verbatim, which is exactly the wire shape under +test. """ from __future__ import annotations import time +from typing import cast import pytest +from anthropic import Anthropic +from anthropic.types import Message, MessageParam, TextBlockParam from pydantic import BaseModel from e2e_config import unique_marker -from e2e_http import Result, unwrap -from endpoints_client import ( - CacheControl, - EndpointsClient, - MessagesResult, - RichMessage, - RichMessagesRequest, - TextBlock, -) from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -48,50 +49,51 @@ CACHE_PRIMING_DEADLINE_SECONDS = 60.0 CACHE_PRIMING_INTERVAL_SECONDS = 3.0 -def _cacheable_system_block(marker: str) -> TextBlock: +def _cacheable_system_block(marker: str) -> TextBlockParam: """A system prompt comfortably above Sonnet's 1024-token minimum cacheable size, unique per run so no other run's cache entry can satisfy the read.""" text = " ".join( f"Reference paragraph {index} for run {marker}." for index in range(300) ) - return TextBlock(text=text, cache_control=CacheControl()) + return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} -def _user_turn(text: str, *, cached: bool = False) -> RichMessage: - block = TextBlock(text=text, cache_control=CacheControl() if cached else None) - return RichMessage(role="user", content=[block]) +def _user_turn(text: str, *, cached: bool = False) -> MessageParam: + block: TextBlockParam = ( + {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + if cached + else {"type": "text", "text": text} + ) + return {"role": "user", "content": [block]} -def _system_reminder_turn() -> RichMessage: - return RichMessage( - role="system", - content=[ - TextBlock( - text="Answer with exactly one word." - ) - ], +def _system_reminder_turn() -> MessageParam: + return cast( + "MessageParam", + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Answer with exactly one word.", + } + ], + }, ) -def _post_messages( - client: EndpointsClient, key: str, body: RichMessagesRequest -) -> Result[MessagesResult]: - return client.proxy.transport.post( - "/v1/messages", - headers=client.proxy.transport.bearer(key), - json=body, - response_type=MessagesResult, - ) +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if block.type == "text") def _register_invoke_deployment( - client: EndpointsClient, resources: ResourceManager, bedrock_model: str + proxy: ProxyClient, resources: ResourceManager, bedrock_model: str ) -> str: model = f"e2e-midsys-{unique_marker()}" - model_id = client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody(model=bedrock_model, aws_region_name=AWS_REGION) ) - resources.defer(lambda: client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model @@ -114,7 +116,7 @@ class PrimedCache(BaseModel): def _prime_prompt_cache( - client: EndpointsClient, key: str, model: str, system_block: TextBlock + client: Anthropic, model: str, system_block: TextBlockParam ) -> PrimedCache: """Send first-turn calls (fresh cache-marked user turn each attempt, identical system prefix) until one both reads the system prefix back from @@ -125,17 +127,19 @@ def _prime_prompt_cache( deadline = time.monotonic() + CACHE_PRIMING_DEADLINE_SECONDS while True: user_text = _first_turn_user_text(unique_marker()) - body = RichMessagesRequest( + usage = client.messages.create( model=model, + max_tokens=64, system=[system_block], messages=[_user_turn(user_text, cached=True)], - ) - usage = unwrap(_post_messages(client, key, body)).usage - if usage.cache_read_input_tokens > 0 and usage.cache_creation_input_tokens > 0: + ).usage + read_tokens = usage.cache_read_input_tokens or 0 + creation_tokens = usage.cache_creation_input_tokens or 0 + if read_tokens > 0 and creation_tokens > 0: return PrimedCache( first_user_text=user_text, - prefix_read_tokens=usage.cache_read_input_tokens, - first_turn_creation_tokens=usage.cache_creation_input_tokens, + prefix_read_tokens=read_tokens, + first_turn_creation_tokens=creation_tokens, ) if time.monotonic() >= deadline: pytest.fail( @@ -151,32 +155,30 @@ class TestBedrockInvokeMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_invoke_deployment( - endpoints_client, resources, FLAGGED_INVOKE_MODEL - ) - key = resources.key(models=[model]) + model = _register_invoke_deployment(proxy, resources, FLAGGED_INVOKE_MODEL) + client = sdk.anthropic(resources.key(models=[model])) system_block = _cacheable_system_block(unique_marker()) - primed = _prime_prompt_cache(endpoints_client, key, model, system_block) + primed = _prime_prompt_cache(client, model, system_block) - reminder_turn_body = RichMessagesRequest( + second = client.messages.create( model=model, + max_tokens=64, system=[system_block], messages=[ _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + {"role": "assistant", "content": [{"type": "text", "text": "OK."}]}, _user_turn("Reply with one word again.", cached=True), ], ) - second = unwrap(_post_messages(endpoints_client, key, reminder_turn_body)) - assert second.text.strip(), ( + assert _text(second).strip(), ( f"{model}: reminder turn returned no completion text" ) - assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + assert (second.usage.cache_read_input_tokens or 0) >= primed.full_prefix_tokens, ( f"{model}: turn with a mid-conversation system reminder read " f"{second.usage.cache_read_input_tokens} cached tokens, expected at " f"least the {primed.full_prefix_tokens} cached on turn one " @@ -191,29 +193,27 @@ class TestBedrockInvokeMidConversationSystem: exercised_on=[], ) def test_unflagged_model_hoists_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_invoke_deployment( - endpoints_client, resources, UNFLAGGED_INVOKE_MODEL - ) - key = resources.key(models=[model]) + model = _register_invoke_deployment(proxy, resources, UNFLAGGED_INVOKE_MODEL) + client = sdk.anthropic(resources.key(models=[model])) - body = RichMessagesRequest( + completion = client.messages.create( model=model, - system=[TextBlock(text="You are terse.")], + max_tokens=64, + system=[{"type": "text", "text": "You are terse."}], messages=[ _user_turn(f"Say hi. Run {unique_marker()}."), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), + {"role": "assistant", "content": [{"type": "text", "text": "Hi."}]}, _user_turn("Say bye."), ], ) - completion = unwrap(_post_messages(endpoints_client, key, body)) assert completion.role == "assistant", ( f"{model}: unexpected role {completion.role!r}" ) - assert completion.text.strip(), ( + assert _text(completion).strip(), ( f"{model}: conversation with a mid-conversation system reminder " f"returned no text; the reminder was forwarded in place to a model " f"that rejects role 'system' inside messages instead of being hoisted" diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py index 97d24e0564b..9a0af34fb5d 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py @@ -24,27 +24,28 @@ entry whose prefix spans ``system`` plus message turns is invalidated when the reminder is hoisted (the ``system`` field mutates and a turn disappears from ``messages``), while an entry ending at the system block itself would survive the hoist and mask the regression. + +Calls go through the real Anthropic SDK (LIT-4577). The SDK's ``MessageParam`` +type only admits user/assistant roles, so the system reminder turn is cast to +it; the SDK serializes the dict verbatim, which is exactly the wire shape under +test. """ from __future__ import annotations import time +from typing import cast import pytest +from anthropic import Anthropic +from anthropic.types import Message, MessageParam, TextBlockParam from pydantic import BaseModel from e2e_config import unique_marker -from e2e_http import Result, unwrap -from endpoints_client import ( - CacheControl, - EndpointsClient, - MessagesResult, - RichMessage, - RichMessagesRequest, - TextBlock, -) from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -68,40 +69,47 @@ def _vertex_params(model: str) -> LiteLLMParamsBody: ) -def _cacheable_system_block(marker: str) -> TextBlock: +def _cacheable_system_block(marker: str) -> TextBlockParam: """A system prompt comfortably above the 1024-token minimum cacheable size, unique per run so no other run's cache entry can satisfy the read.""" text = " ".join(f"Reference paragraph {index} for run {marker}." for index in range(300)) - return TextBlock(text=text, cache_control=CacheControl()) + return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} -def _user_turn(text: str, *, cached: bool = False) -> RichMessage: - block = TextBlock(text=text, cache_control=CacheControl() if cached else None) - return RichMessage(role="user", content=[block]) +def _user_turn(text: str, *, cached: bool = False) -> MessageParam: + block: TextBlockParam = ( + {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + if cached + else {"type": "text", "text": text} + ) + return {"role": "user", "content": [block]} -def _system_reminder_turn() -> RichMessage: - return RichMessage( - role="system", - content=[TextBlock(text="Answer with exactly one word.")], +def _system_reminder_turn() -> MessageParam: + return cast( + "MessageParam", + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Answer with exactly one word.", + } + ], + }, ) -def _post_messages(client: EndpointsClient, key: str, body: RichMessagesRequest) -> Result[MessagesResult]: - return client.proxy.transport.post( - "/v1/messages", - headers=client.proxy.transport.bearer(key), - json=body, - response_type=MessagesResult, - ) +def _text(message: Message) -> str: + return "".join(block.text for block in message.content if block.type == "text") def _register_deployment( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody + proxy: ProxyClient, resources: ResourceManager, params: LiteLLMParamsBody ) -> str: model = f"e2e-midsys-{unique_marker()}" - model_id = client.create_model(model, params) - resources.defer(lambda: client.delete_model(model_id)) + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) return model @@ -124,7 +132,7 @@ class PrimedCache(BaseModel): def _prime_prompt_cache( - client: EndpointsClient, key: str, model: str, system_block: TextBlock + client: Anthropic, model: str, system_block: TextBlockParam ) -> PrimedCache: """Send first-turn calls (fresh cache-marked user turn each attempt, identical system prefix) until one both reads the system prefix back from @@ -135,17 +143,19 @@ def _prime_prompt_cache( deadline = time.monotonic() + CACHE_PRIMING_DEADLINE_SECONDS while True: user_text = _first_turn_user_text(unique_marker()) - body = RichMessagesRequest( + usage = client.messages.create( model=model, + max_tokens=64, system=[system_block], messages=[_user_turn(user_text, cached=True)], - ) - usage = unwrap(_post_messages(client, key, body)).usage - if usage.cache_read_input_tokens > 0 and usage.cache_creation_input_tokens > 0: + ).usage + read_tokens = usage.cache_read_input_tokens or 0 + creation_tokens = usage.cache_creation_input_tokens or 0 + if read_tokens > 0 and creation_tokens > 0: return PrimedCache( first_user_text=user_text, - prefix_read_tokens=usage.cache_read_input_tokens, - first_turn_creation_tokens=usage.cache_creation_input_tokens, + prefix_read_tokens=read_tokens, + first_turn_creation_tokens=creation_tokens, ) if time.monotonic() >= deadline: pytest.fail( @@ -156,28 +166,31 @@ def _prime_prompt_cache( def _assert_flagged_model_keeps_cache( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + params: LiteLLMParamsBody, ) -> None: - model = _register_deployment(client, resources, params) - key = resources.key(models=[model]) + model = _register_deployment(proxy, resources, params) + client = sdk.anthropic(resources.key(models=[model])) system_block = _cacheable_system_block(unique_marker()) - primed = _prime_prompt_cache(client, key, model, system_block) + primed = _prime_prompt_cache(client, model, system_block) - reminder_turn_body = RichMessagesRequest( + second = client.messages.create( model=model, + max_tokens=64, system=[system_block], messages=[ _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + {"role": "assistant", "content": [{"type": "text", "text": "OK."}]}, _user_turn("Reply with one word again.", cached=True), ], ) - second = unwrap(_post_messages(client, key, reminder_turn_body)) - assert second.text.strip(), f"{model}: reminder turn returned no completion text" - assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + assert _text(second).strip(), f"{model}: reminder turn returned no completion text" + assert (second.usage.cache_read_input_tokens or 0) >= primed.full_prefix_tokens, ( f"{model}: turn with a mid-conversation system reminder read " f"{second.usage.cache_read_input_tokens} cached tokens, expected at " f"least the {primed.full_prefix_tokens} cached on turn one " @@ -189,25 +202,28 @@ def _assert_flagged_model_keeps_cache( def _assert_unflagged_model_hoists_and_succeeds( - client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + params: LiteLLMParamsBody, ) -> None: - model = _register_deployment(client, resources, params) - key = resources.key(models=[model]) + model = _register_deployment(proxy, resources, params) + client = sdk.anthropic(resources.key(models=[model])) - body = RichMessagesRequest( + completion = client.messages.create( model=model, - system=[TextBlock(text="You are terse.")], + max_tokens=64, + system=[{"type": "text", "text": "You are terse."}], messages=[ _user_turn(f"Say hi. Run {unique_marker()}."), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), + {"role": "assistant", "content": [{"type": "text", "text": "Hi."}]}, _user_turn("Say bye."), ], ) - completion = unwrap(_post_messages(client, key, body)) assert completion.role == "assistant", f"{model}: unexpected role {completion.role!r}" - assert completion.text.strip(), ( + assert _text(completion).strip(), ( f"{model}: conversation with a mid-conversation system reminder returned " f"no text; the reminder was forwarded in place to a model that rejects " f"role 'system' inside messages instead of being hoisted" @@ -223,19 +239,19 @@ class TestAzureFoundryMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - _assert_flagged_model_keeps_cache(endpoints_client, resources, _azure_params(self.FLAGGED_MODEL)) + _assert_flagged_model_keeps_cache(proxy, resources, sdk, _azure_params(self.FLAGGED_MODEL)) @pytest.mark.covers( "llm.messages.azure_foundry.mid_conversation_system.nonstream.works", exercised_on=[], ) def test_unflagged_model_hoists_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: _assert_unflagged_model_hoists_and_succeeds( - endpoints_client, resources, _azure_params(self.UNFLAGGED_MODEL) + proxy, resources, sdk, _azure_params(self.UNFLAGGED_MODEL) ) @@ -248,17 +264,17 @@ class TestVertexMidConversationSystem: exercised_on=[], ) def test_flagged_model_keeps_prompt_cache_across_system_reminder( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - _assert_flagged_model_keeps_cache(endpoints_client, resources, _vertex_params(self.FLAGGED_MODEL)) + _assert_flagged_model_keeps_cache(proxy, resources, sdk, _vertex_params(self.FLAGGED_MODEL)) @pytest.mark.covers( "llm.messages.vertex.mid_conversation_system.nonstream.works", exercised_on=[], ) def test_unflagged_model_hoists_system_reminder_and_succeeds( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: _assert_unflagged_model_hoists_and_succeeds( - endpoints_client, resources, _vertex_params(self.UNFLAGGED_MODEL) + proxy, resources, sdk, _vertex_params(self.UNFLAGGED_MODEL) ) diff --git a/tests/e2e/llm_translation/test_moderations_e2e.py b/tests/e2e/llm_translation/test_moderations_e2e.py index 69cf4414a48..bc26d79dcab 100644 --- a/tests/e2e/llm_translation/test_moderations_e2e.py +++ b/tests/e2e/llm_translation/test_moderations_e2e.py @@ -1,19 +1,22 @@ """Live e2e: POST /v1/moderations classifies content against the provider policy. -Registers OpenAI's omni moderation model at runtime and asserts the product -promise on both sides of the decision: clearly violent text comes back flagged -with at least one policy category tripped, and benign text comes back not flagged. +Registers OpenAI's omni moderation model at runtime, drives it through the real +OpenAI SDK (LIT-4577), and asserts the product promise on both sides of the +decision: clearly violent text comes back flagged with at least one policy +category tripped, and benign text comes back not flagged. """ from __future__ import annotations import pytest +from openai.types import Moderation +from pydantic import TypeAdapter from e2e_config import unique_marker -from e2e_http import unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e @@ -21,45 +24,49 @@ VIOLENT_TEXT = "I am going to find you and kill you, and I will hurt everyone yo BENIGN_TEXT = "I enjoyed the sunny afternoon and a relaxing walk in the park today." -def _register_moderation_model( - endpoints_client: EndpointsClient, resources: ResourceManager -) -> str: +def _register_moderation_model(proxy: ProxyClient, resources: ResourceManager) -> str: model = f"e2e-moderation-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="openai/omni-moderation-latest", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) return model +_CATEGORY_FLAGS = TypeAdapter(dict[str, bool | None]) + + +def _flagged_categories(item: Moderation) -> tuple[str, ...]: + flags = _CATEGORY_FLAGS.validate_python(item.categories.model_dump()) + return tuple(name for name, hit in flags.items() if hit) + + class TestModerations: @pytest.mark.covers("llm.moderations.openai.basic.nonstream.works") def test_moderations_flags_violent_content( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_moderation_model(endpoints_client, resources) - key = resources.key() + model = _register_moderation_model(proxy, resources) + client = sdk.openai(resources.key()) - result = unwrap(endpoints_client.moderations(key, model, VIOLENT_TEXT)) - item = result.first - assert item is not None, f"/moderations returned no results: {result}" - assert item.flagged, f"violent text was not flagged: {item}" - assert item.flagged_categories, ( - f"flagged result reported no true category: {item}" - ) + moderation = client.moderations.create(model=model, input=VIOLENT_TEXT) + assert moderation.results, f"/moderations returned no results: {moderation!r}" + item = moderation.results[0] + assert item.flagged, f"violent text was not flagged: {item!r}" + assert _flagged_categories(item), f"flagged result reported no true category: {item!r}" def test_moderations_passes_benign_content( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = _register_moderation_model(endpoints_client, resources) - key = resources.key() + model = _register_moderation_model(proxy, resources) + client = sdk.openai(resources.key()) - result = unwrap(endpoints_client.moderations(key, model, BENIGN_TEXT)) - item = result.first - assert item is not None, f"/moderations returned no results: {result}" + moderation = client.moderations.create(model=model, input=BENIGN_TEXT) + assert moderation.results, f"/moderations returned no results: {moderation!r}" + item = moderation.results[0] assert not item.flagged, ( - f"benign text was flagged as {item.flagged_categories}: {item}" + f"benign text was flagged as {_flagged_categories(item)}: {item!r}" ) diff --git a/tests/e2e/llm_translation/test_ocr_rust_e2e.py b/tests/e2e/llm_translation/test_ocr_rust_e2e.py index cdbf1883314..986ac675e98 100644 --- a/tests/e2e/llm_translation/test_ocr_rust_e2e.py +++ b/tests/e2e/llm_translation/test_ocr_rust_e2e.py @@ -22,9 +22,9 @@ import pytest from e2e_config import unique_marker from e2e_http import unwrap -from endpoints_client import EndpointsClient from lifecycle import ResourceManager from models import LiteLLMParamsBody, OcrBody, OcrDocument, OcrResponse +from proxy_client import ProxyClient pytestmark = pytest.mark.e2e @@ -143,14 +143,14 @@ def _assert_ocr_document(response: OcrResponse) -> None: class TestRustOcrGateway: @pytest.mark.parametrize("case", RUST_OCR_CASES, ids=_CASE_IDS) def test_rust_ocr_response( - self, endpoints_client: EndpointsClient, resources: ResourceManager, case: _OcrCase + self, proxy: ProxyClient, resources: ResourceManager, case: _OcrCase ) -> None: model = f"rust-ocr-{case.suffix}-{unique_marker()}" - model_id = endpoints_client.create_model(model, case.provider.litellm_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + model_id = proxy.create_model(model, case.provider.litellm_params()) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - response = unwrap(endpoints_client.proxy.ocr(key, OcrBody(model=model, document=case.document))) + response = unwrap(proxy.ocr(key, OcrBody(model=model, document=case.document))) _assert_ocr_document(response) diff --git a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py index 045988334d5..00ca810db15 100644 --- a/tests/e2e/llm_translation/test_passthrough_headers_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_headers_e2e.py @@ -18,9 +18,8 @@ from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import AuthHeaders, NoBody, require_successful_call, unwrap -from endpoints_client import MessagesResult from lifecycle import ResourceManager -from models import ChatMessage, KeyGenerateBody +from models import AnthropicMessagesResponse, ChatMessage, KeyGenerateBody from passthrough_client import PassthroughClient pytestmark = pytest.mark.e2e @@ -128,8 +127,9 @@ class TestPassthroughHeaders: json=_messages_body(), ) require_successful_call(result) - completion = MessagesResult.model_validate_json(result.body) - assert completion.text.strip(), ( + completion = AnthropicMessagesResponse.model_validate_json(result.body) + text = "".join(block.text or "" for block in (completion.content or [])) + assert text.strip(), ( f"static x-api-key must reach Anthropic for the call to succeed at all; got {result.body[:300]}" ) diff --git a/tests/e2e/llm_translation/test_rerank_e2e.py b/tests/e2e/llm_translation/test_rerank_e2e.py index 0857ff65a52..c0e4d7def42 100644 --- a/tests/e2e/llm_translation/test_rerank_e2e.py +++ b/tests/e2e/llm_translation/test_rerank_e2e.py @@ -1,8 +1,9 @@ """Live e2e: POST /v1/rerank ranks documents by relevance. Registers a Cohere rerank deployment at runtime and asserts the endpoint returns -scored results within the requested top_n. Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +scored results within the requested top_n. No official OpenAI/Anthropic SDK +covers /v1/rerank, so the call rides the shared typed transport via +ProxyClient.rerank. """ from __future__ import annotations @@ -10,10 +11,10 @@ from __future__ import annotations import pytest from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, RerankResult +from e2e_http import unwrap from lifecycle import ResourceManager -from models import LiteLLMParamsBody +from models import LiteLLMParamsBody, RerankBody, RerankResponse +from proxy_client import ProxyClient pytestmark = pytest.mark.e2e @@ -26,39 +27,39 @@ DOCUMENTS = [ QUERY = "What is the capital of the United States?" -def _assert_top_n_scored(body: str) -> None: - parsed = RerankResult.model_validate_json(body) - assert parsed.results, f"/rerank returned no results: {body[:300]}" - assert len(parsed.results) <= 3, f"top_n=3 not honored: {body[:300]}" - assert parsed.results[0].relevance_score is not None, ( - f"top rerank result has no relevance_score: {body[:300]}" +def _assert_top_n_scored(response: RerankResponse) -> None: + assert response.results, f"/rerank returned no results: {response!r}" + assert len(response.results) <= 3, f"top_n=3 not honored: {response!r}" + assert response.results[0].relevance_score is not None, ( + f"top rerank result has no relevance_score: {response!r}" ) class TestRerank: @pytest.mark.covers("llm.rerank.cohere.basic.nonstream.works") def test_rerank_scores_top_n( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: model = f"e2e-rerank-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody(model="cohere/rerank-v3.5", api_key="os.environ/COHERE_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) - require_successful_call(result) - _assert_top_n_scored(result.body) + response = unwrap( + proxy.rerank(key, RerankBody(model=model, query=QUERY, documents=DOCUMENTS, top_n=3)) + ) + _assert_top_n_scored(response) @pytest.mark.covers("llm.rerank.bedrock.basic.nonstream.works", exercised_on=["rerank"]) def test_bedrock_rerank_scores_top_n( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager ) -> None: require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") model = f"e2e-bedrock-rerank-{unique_marker()}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="bedrock/amazon.rerank-v1:0", @@ -67,9 +68,10 @@ class TestRerank: aws_region_name="os.environ/AWS_REGION", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) + resources.defer(lambda: proxy.delete_model(model_id)) key = resources.key() - result = endpoints_client.rerank(key, model, QUERY, DOCUMENTS, top_n=3) - require_successful_call(result) - _assert_top_n_scored(result.body) + response = unwrap( + proxy.rerank(key, RerankBody(model=model, query=QUERY, documents=DOCUMENTS, top_n=3)) + ) + _assert_top_n_scored(response) diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index d24d2b53b71..93f844505d0 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -1,8 +1,8 @@ """Live e2e: POST /v1/responses returns a real completion. -Registers an OpenAI deployment at runtime, drives the Responses API through the -gateway, and asserts output text came back. Migrated from -litellm-regression-tests/tests/test_inference_endpoints.py. +Registers an OpenAI deployment at runtime and drives the Responses API through +the gateway with the real OpenAI SDK, the client customers actually use +(LIT-4577), asserting output text came back. """ from __future__ import annotations @@ -11,34 +11,45 @@ import json from typing import cast import pytest -from pydantic import BaseModel, ValidationError +from openai.types.responses import ( + FunctionToolParam, + Response, + ResponseFunctionToolCall, + ResponseInputParam, +) +from pydantic import BaseModel from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call -from endpoints_client import ( - EndpointsClient, - FunctionParameterProperty, - FunctionParameters, - ResponsesFunctionTool, - ResponsesOutputTextDeltaEvent, - ResponsesResult, - ResponsesStreamEventType, -) from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +INSTRUCTIONS = "You are a helpful assistant" +CAT_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg" -WEATHER_TOOL = ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), -) +WEATHER_TOOL: FunctionToolParam = { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + "strict": False, +} + + +def _openai_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY") + + +def _anthropic_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY") def _bedrock_params() -> LiteLLMParamsBody: @@ -50,6 +61,25 @@ def _bedrock_params() -> LiteLLMParamsBody: ) +def _register(proxy: ProxyClient, resources: ResourceManager, params: LiteLLMParamsBody) -> str: + model = f"e2e-responses-{unique_marker()}" + model_id = proxy.create_model(model, params) + resources.defer(lambda: proxy.delete_model(model_id)) + return model + + +def _function_calls(response: Response) -> tuple[ResponseFunctionToolCall, ...]: + return tuple(item for item in response.output if isinstance(item, ResponseFunctionToolCall)) + + +def _assert_weather_call(response: Response) -> None: + function_call = next((call for call in _function_calls(response) if call.name == "get_weather"), None) + assert function_call is not None, f"no get_weather function call: {response.output!r}" + raw_arguments = cast(object, json.loads(function_call.arguments)) + arguments = WeatherArguments.model_validate(raw_arguments) + assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + + class WeatherArguments(BaseModel): location: str @@ -57,242 +87,157 @@ class WeatherArguments(BaseModel): class TestResponses: @pytest.mark.covers("llm.responses.openai.basic.nonstream.works") def test_responses_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS + ) + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" @pytest.mark.covers("llm.responses.openai.basic.stream.works") def test_responses_streaming_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word", stream=True) - require_successful_call(result) - delta_events = tuple( - parsed - for event in result.stream_events - if (parsed := _parse_stream_event(event)) is not None + stream = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS, stream=True + ) + events = list(stream) + assert events, "responses stream returned no events" + deltas = [event.delta for event in events if event.type == "response.output_text.delta"] + assert any(delta for delta in deltas), "responses stream returned no text deltas" + assert events[-1].type == "response.completed", ( + f"responses stream did not terminate with response.completed: {events[-1].type}" ) - - assert any(event.delta for event in delta_events), "responses stream returned no text deltas" - assert result.stream_events, "responses stream returned no events" - assert ( - ResponsesStreamEventType.model_validate_json(result.stream_events[-1]).type - == "response.completed" - ), "responses stream did not terminate with response.completed" @pytest.mark.covers("llm.responses.openai.basic.nonstream.cost_logged") def test_responses_logs_cost( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) + + raw = client.responses.with_raw_response.create( + model=model, input=f"reply with one word {unique_marker()}", instructions=INSTRUCTIONS + ) + response = raw.parse() + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" + assert raw.headers.get("x-litellm-call-id") and response.id, ( + f"missing response identifiers: id={response.id!r}, headers={dict(raw.headers)}" ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() - result = endpoints_client.responses(key, model, f"reply with one word {unique_marker()}") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" - assert result.call_id and parsed.id, f"missing response identifiers: {result.body[:300]}" - - rows = endpoints_client.proxy.poll_logs_for_request_id( - parsed.id, + rows = proxy.poll_logs_for_request_id( + response.id, predicate=lambda logged_rows: any((row.spend or 0) > 0 for row in logged_rows), ) row = next((logged_row for logged_row in rows if (logged_row.spend or 0) > 0), None) - assert row is not None, f"no costed spend row for response id {parsed.id}" + assert row is not None, f"no costed spend row for response id {response.id}" assert "gpt-4o-mini" in (row.model or ""), f"unexpected spend row model: {row.model}" @pytest.mark.covers("llm.responses.openai.tool_use.nonstream.works") def test_responses_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _openai_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, - model, - "What is the weather in San Francisco? Use the get_weather tool.", - [ - ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), - ) - ], + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next( - (call for call in parsed.function_calls if call.name == "get_weather"), - None, - ) - assert function_call is not None, f"no get_weather function call: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + _assert_weather_call(response) @pytest.mark.covers("llm.responses.openai.vision.nonstream.works") def test_responses_vision_describes_image( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, + model = _register( + proxy, + resources, LiteLLMParamsBody(model="openai/gpt-4o", api_key="os.environ/OPENAI_API_KEY"), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + client = sdk.openai(resources.key()) - result = endpoints_client.responses_vision( - key, - model, - "What animal is shown in this image? Answer in one word", - "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg", + vision_input: ResponseInputParam = [ + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "What animal is shown in this image? Answer in one word", + }, + {"type": "input_image", "image_url": CAT_IMAGE_URL, "detail": "auto"}, + ], + } + ] + response = client.responses.create(model=model, input=vision_input, instructions=INSTRUCTIONS) + text = response.output_text.strip().lower() + assert text, f"/responses vision returned no output text: {response.output!r}" + assert any(keyword in text for keyword in ("cat", "feline")), ( + f"vision response did not describe the image: {text[:300]}" ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - text = parsed.text.strip().lower() - assert text, f"/responses vision returned no output text: {result.body[:300]}" - assert any( - keyword in text - for keyword in ("cat", "feline") - ), f"vision response did not describe the image: {parsed.text[:300]}" @pytest.mark.covers("llm.responses.anthropic.basic.nonstream.works") def test_responses_anthropic_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _anthropic_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS + ) + assert response.output_text.strip(), f"/responses returned no output text: {response.output!r}" @pytest.mark.covers("llm.responses.anthropic.tool_use.nonstream.works") def test_responses_anthropic_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model="anthropic/claude-haiku-4-5", api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _anthropic_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, - model, - "What is the weather in San Francisco? Use the get_weather tool.", - [ - ResponsesFunctionTool( - name="get_weather", - description="Get the weather for a location", - parameters=FunctionParameters( - properties={"location": FunctionParameterProperty(type="string")}, - required=["location"], - ), - ) - ], + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next( - (call for call in parsed.function_calls if call.name == "get_weather"), - None, - ) - assert function_call is not None, f"no get_weather function call: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" + _assert_weather_call(response) @pytest.mark.covers("llm.responses.bedrock_converse.basic.nonstream.works") def test_responses_bedrock_returns_completion( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model(model, _bedrock_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _bedrock_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses(key, model, "reply with one word") - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - assert parsed.text.strip(), f"/responses over bedrock returned no output text: {result.body[:300]}" + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS + ) + assert response.output_text.strip(), ( + f"/responses over bedrock returned no output text: {response.output!r}" + ) @pytest.mark.covers("llm.responses.bedrock_converse.tool_use.nonstream.works") def test_responses_bedrock_returns_function_call( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: require_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION") - model = f"e2e-responses-{unique_marker()}" - model_id = endpoints_client.create_model(model, _bedrock_params()) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + model = _register(proxy, resources, _bedrock_params()) + client = sdk.openai(resources.key()) - result = endpoints_client.responses_with_tools( - key, model, "What is the weather in San Francisco? Use the get_weather tool.", [WEATHER_TOOL] + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], ) - require_successful_call(result) - parsed = ResponsesResult.model_validate_json(result.body) - function_call = next((call for call in parsed.function_calls if call.name == "get_weather"), None) - assert function_call is not None, f"no get_weather function call over bedrock: {result.body[:500]}" - assert function_call.arguments is not None - raw_arguments = cast(object, json.loads(function_call.arguments)) - arguments = WeatherArguments.model_validate(raw_arguments) - assert arguments.location, f"function call arguments missing location: {function_call.arguments}" - - -def _parse_stream_event( - event: str, -) -> ResponsesOutputTextDeltaEvent | None: - try: - return ResponsesOutputTextDeltaEvent.model_validate_json(event) - except ValidationError: - return None + _assert_weather_call(response) diff --git a/tests/e2e/llm_translation/test_responses_metadata_e2e.py b/tests/e2e/llm_translation/test_responses_metadata_e2e.py index 6cf24348095..87727bde4d8 100644 --- a/tests/e2e/llm_translation/test_responses_metadata_e2e.py +++ b/tests/e2e/llm_translation/test_responses_metadata_e2e.py @@ -1,8 +1,8 @@ """Live e2e: /v1/responses with store + metadata (LIT-1201 customer path). -Customers attach metadata and store=true, then continue with previous_response_id. -Both turns must succeed, and any Redis keys written for the session must carry a -positive TTL (not unbounded). +Customers attach metadata and store=true through the OpenAI SDK, then continue +with previous_response_id. Both turns must succeed, and any Redis keys written +for the session must carry a positive TTL (not unbounded). """ from __future__ import annotations @@ -15,21 +15,14 @@ import pytest from pydantic import BaseModel, ConfigDict from e2e_config import require_env, unique_marker -from e2e_http import require_successful_call -from endpoints_client import EndpointsClient, ResponsesResult from lifecycle import ResourceManager from models import LiteLLMParamsBody +from proxy_client import ProxyClient +from sdk_clients import SdkClients pytestmark = pytest.mark.e2e - -class ResponsesMetadataBody(BaseModel): - model: str - input: str - store: bool = True - metadata: dict[str, str] - previous_response_id: str | None = None - instructions: str | None = "You are a helpful assistant." +INSTRUCTIONS = "You are a helpful assistant." class RedisKeyInfo(BaseModel): @@ -67,50 +60,42 @@ class TestResponsesMetadata: exercised_on=["responses"], ) def test_store_metadata_continues_and_redis_keys_have_ttl( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: # Anthropic avoids OpenAI/Gemini quota flakes; Responses translation still # exercises store + metadata + previous_response_id on the proxy. marker = unique_marker() model = f"e2e-resp-meta-{marker}" - model_id = endpoints_client.create_model( + model_id = proxy.create_model( model, LiteLLMParamsBody( model="anthropic/claude-haiku-4-5-20251001", api_key="os.environ/ANTHROPIC_API_KEY", ), ) - resources.defer(lambda: endpoints_client.delete_model(model_id)) - key = resources.key() + resources.defer(lambda: proxy.delete_model(model_id)) + client = sdk.openai(resources.key()) - first = endpoints_client.proxy.transport.send( - "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), - json=ResponsesMetadataBody( - model=model, - input=f"Remember marker {marker}. Reply with one word.", - metadata={"session_id": marker, "customer": "e2e"}, - ), + first = client.responses.create( + model=model, + input=f"Remember marker {marker}. Reply with one word.", + store=True, + metadata={"session_id": marker, "customer": "e2e"}, + instructions=INSTRUCTIONS, ) - require_successful_call(first) - parsed = ResponsesResult.model_validate_json(first.body) - assert parsed.id, f"responses must return an id: {first.body[:300]}" - assert parsed.text.strip(), f"responses returned empty text: {first.body[:300]}" + assert first.id, f"responses must return an id: {first!r}" + assert first.output_text.strip(), f"responses returned empty text: {first.output!r}" - second = endpoints_client.proxy.transport.send( - "/v1/responses", - headers=endpoints_client.proxy.transport.bearer(key), - json=ResponsesMetadataBody( - model=model, - input="Reply with the single word ok.", - previous_response_id=parsed.id, - metadata={"session_id": marker, "turn": "2"}, - ), + second = client.responses.create( + model=model, + input="Reply with the single word ok.", + store=True, + previous_response_id=first.id, + metadata={"session_id": marker, "turn": "2"}, + instructions=INSTRUCTIONS, ) - require_successful_call(second) - second_parsed = ResponsesResult.model_validate_json(second.body) - assert second_parsed.text.strip(), ( - f"previous_response_id follow-up returned empty text: {second.body[:300]}" + assert second.output_text.strip(), ( + f"previous_response_id follow-up returned empty text: {second.output!r}" ) time.sleep(1.0) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index b3ea9346180..722722f4e19 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -364,6 +364,25 @@ class EmbedResponse(BaseModel): model: str | None = None +# ---------- rerank ---------- + + +class RerankBody(BaseModel): + model: str + query: str + documents: list[str] + top_n: int + + +class RerankItem(BaseModel): + index: int | None = None + relevance_score: float | None = None + + +class RerankResponse(BaseModel): + results: list[RerankItem] = [] + + # ---------- ocr ---------- diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..234c37d2f77 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -57,6 +57,8 @@ from models import ( ModelUpdateBody, OcrBody, OcrResponse, + RerankBody, + RerankResponse, SpendLogRow, SpendLogs, SpendLogsPage, @@ -291,6 +293,16 @@ class ProxyClient: response_type=OcrResponse, ) + def rerank(self, key: str, body: RerankBody) -> Result[RerankResponse]: + """POST /v1/rerank (Cohere-format). No official OpenAI/Anthropic SDK + covers this route, so it stays on the shared typed transport.""" + return self.transport.post( + "/v1/rerank", + headers=self.transport.bearer(key), + json=body, + response_type=RerankResponse, + ) + def count_tokens(self, key: str, body: CountTokensBody) -> Result[CountTokensResponse]: """POST /v1/messages/count_tokens (Anthropic-native). Sends the anthropic-version header so the native path accepts it; harmless on the diff --git a/uv.lock b/uv.lock index b3d6fccff26..caf976b9202 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-19T00:00:06.091071Z" +exclude-newer = "2026-07-20T03:20:37.107777782Z" exclude-newer-span = "P3D" [manifest] @@ -4300,6 +4300,7 @@ dev = [ { name = "vcrpy" }, ] e2e-dev = [ + { name = "anthropic" }, { name = "locust" }, { name = "playwright" }, { name = "websockets" }, @@ -4477,6 +4478,7 @@ dev = [ { name = "vcrpy", specifier = "==8.2.1" }, ] e2e-dev = [ + { name = "anthropic", specifier = "==0.84.0" }, { name = "locust", specifier = "==2.45.0" }, { name = "playwright", specifier = "==1.61.0" }, { name = "websockets", specifier = ">=15.0.1,<16.0" }, From 0e34c18d7ed9a846258fe6456117b8143134bfc6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 23 Jul 2026 06:32:29 +0000 Subject: [PATCH 002/246] docs(e2e): name the e2e-dev group in the documented suite run commands The llm_translation suite needs the e2e-dev dependency group at collection time (websockets for the realtime folder, now also the anthropic SDK for the sdk fixture). make bootstrap installs the group, but the documented pytest command did not name it, so a default dev-group environment failed collection. Naming the group on uv run makes the command work from any environment state --- tests/e2e/CONTRIBUTING.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index fa0174f686c..e542c123988 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -35,10 +35,10 @@ The suites run against a live proxy, so bring one up first by running the litell curl -fs http://localhost:4000/health/liveliness ``` -4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`): +4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`). The suites' client dependencies (the provider SDKs, websockets) live in the `e2e-dev` dependency group; `make bootstrap` installs it, and naming the group on the run keeps the command working from any environment state: ```bash - uv run pytest tests/e2e/llm_translation/ -v + uv run --group e2e-dev pytest tests/e2e/llm_translation/ -v ``` The browser tests in the `management/` suite drive the dashboard the proxy serves at `/ui` through playwright, an optional dependency behind `importorskip` (the suite's API tests run without it). It lives in the `e2e-dev` dependency group; install it along with its browser: @@ -148,7 +148,7 @@ Before you push ```bash litellm --config .yml --port 4000 - uv run pytest tests/e2e// -v + uv run --group e2e-dev pytest tests/e2e// -v ``` 4. Capture screenshots of the test run and attach them to the PR as proof From 6bbcf1dbbf3d533ee1948b9b3183d8962ace4263 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:08:08 +0000 Subject: [PATCH 003/246] fix(mcp): keep the streamable-HTTP routing peek on a UTF-8 boundary Fixes https://github.com/BerriAI/litellm/issues/34917 --- .../proxy/_experimental/mcp_server/server.py | 24 +++- .../mcp_server/test_mcp_server.py | 104 ++++++++++++++++++ 2 files changed, 125 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 14673cf12c1..af5da275961 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -237,6 +237,24 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: return False +def _utf8_boundary_prefix(data: bytes) -> bytes: + """``data`` with any trailing incomplete UTF-8 sequence removed. + + Cutting a body at a fixed byte budget can land in the middle of a multibyte + character, and ``json.loads`` on such bytes raises ``UnicodeDecodeError`` + rather than ``JSONDecodeError``. Trimming to a character boundary keeps the + truncated peek decodable so callers only have to handle malformed JSON. + """ + for trailing in range(0, min(3, len(data)) + 1): + candidate = data[: len(data) - trailing] + try: + candidate.decode("utf-8") + except UnicodeDecodeError: + continue + return candidate + return data + + def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None: """The W3C trace context (``traceparent``/``tracestate``) the MCP client propagated in the request's ``params._meta`` (SEP-414), or ``None``. @@ -3411,7 +3429,7 @@ if MCP_AVAILABLE: try: data = json.loads(body) return isinstance(data, dict) and data.get("method") == "initialize" - except (json.JSONDecodeError, TypeError): + except (json.JSONDecodeError, UnicodeDecodeError, TypeError): return False async def _read_request_body_for_routing( @@ -3462,7 +3480,7 @@ if MCP_AVAILABLE: # directly from the original `receive` via wrapped_receive. break - return consumed_messages, b"".join(body_chunks) + return consumed_messages, _utf8_boundary_prefix(b"".join(body_chunks)) async def _handle_stale_mcp_session( scope: Scope, @@ -4227,7 +4245,7 @@ if MCP_AVAILABLE: "MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock", _peeked.get("id"), ) - except (json.JSONDecodeError, TypeError): + except (json.JSONDecodeError, UnicodeDecodeError, TypeError): # Peek cap truncated the body, so it can't be fully parsed. # Scan the top-level keys (depth-aware) instead of a flat # substring search: a response's result payload may nest a diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 1753b0d92a8..7f79e5aebda 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,5 +1,6 @@ import asyncio import contextvars +import json from datetime import datetime, timedelta from unittest.mock import AsyncMock, MagicMock, patch @@ -1689,6 +1690,109 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body(): assert total_streamed == len(first_chunk) + sum(len(b) for b in oversized_tail) +@pytest.mark.asyncio +async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap(): + """ + A tool-call POST whose UTF-8 body is larger than the routing peek cap, with a + multibyte character straddling the cap boundary, must still be forwarded + intact instead of blowing up with a UnicodeDecodeError 500. + + Regression test for https://github.com/BerriAI/litellm/issues/34917 + """ + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + peek_cap = mcp_server._MCP_ROUTING_PEEK_MAX_BYTES + + def _splits_multibyte_at_cap(candidate: bytes) -> bool: + try: + candidate[:peek_cap].decode("utf-8") + except UnicodeDecodeError: + return True + return False + + def _build_body() -> bytes: + for pad in range(4): + candidate = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "update_full_document" + "x" * pad, + "arguments": {"markdown": "щ" * 3000}, + }, + }, + ensure_ascii=False, + ).encode("utf-8") + if len(candidate) > peek_cap and _splits_multibyte_at_cap(candidate): + return candidate + raise AssertionError("could not build a body splitting a multibyte char at the peek cap") + + body = _build_body() + + messages = [{"type": "http.request", "body": body, "more_body": False}] + receive_calls = {"count": 0} + + async def receive(): + idx = receive_calls["count"] + receive_calls["count"] += 1 + return messages[idx] + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/progress_test", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer test-key"), + ], + } + send = AsyncMock() + + streamed_chunks = [] + + async def stateless_handle(s, r, se): + while True: + msg = await r() + if msg.get("type") != "http.request": + break + streamed_chunks.append(msg.get("body", b"") or b"") + if not msg.get("more_body", False): + break + + async def stateful_handle(s, r, se): + raise AssertionError("non-initialize POST should not reach stateful manager") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(MagicMock(), None, ["progress_test"], None, None, None), + ), + patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object(session_manager_stateless, "handle_request", side_effect=stateless_handle), + patch.object(session_manager_stateful, "handle_request", side_effect=stateful_handle), + patch.object(session_manager_stateless, "_server_instances", {}), + patch.object(session_manager_stateful, "_server_instances", {}), + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert send.await_count == 0, f"unexpected response emitted by the proxy: {send.await_args_list}" + assert b"".join(streamed_chunks) == body + + @pytest.mark.asyncio async def test_enforce_stateful_session_cap_evicts_oldest_idle_then_rejects(): """ From 4a6a387ca1e511e35858fee0c92fe3e3415d03ee Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:48:22 +0000 Subject: [PATCH 004/246] fix(mcp): follow nextCursor on paginated tools/prompts/resources list operations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/experimental_mcp_client/client.py | 56 ++++++----- litellm/experimental_mcp_client/pagination.py | 92 +++++++++++++++++++ litellm/experimental_mcp_client/tools.py | 7 +- .../mcp_server/rest_endpoints.py | 6 +- .../test_mcp_client.py | 52 ++++++++++- .../test_pagination.py | 80 ++++++++++++++++ 7 files changed, 258 insertions(+), 36 deletions(-) create mode 100644 litellm/experimental_mcp_client/pagination.py create mode 100644 tests/test_litellm/experimental_mcp_client/test_pagination.py diff --git a/litellm/constants.py b/litellm/constants.py index 9a50797f517..11f35177636 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -136,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0" MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) +MCP_LIST_MAX_PAGES: Final = int(os.getenv("LITELLM_MCP_LIST_MAX_PAGES", "100")) # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index f0a1bff8fdc..814074d35a4 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -48,6 +48,12 @@ from pydantic import AnyUrl from litellm._logging import verbose_logger from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR +from litellm.experimental_mcp_client.pagination import ( + list_all_prompts, + list_all_resource_templates, + list_all_resources, + list_all_tools, +) from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( @@ -603,17 +609,17 @@ class MCPClient: """ verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_tools_operation(session: ClientSession): - return await session.list_tools() + async def _list_tools_operation(session: ClientSession) -> tuple[MCPTool, ...]: + return await list_all_tools(session, self.server_url or "stdio") try: - result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) - tool_count: Final = len(result.tools) - tool_names: Final = [tool.name for tool in result.tools] + tools: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) + tool_count: Final = len(tools) + tool_names: Final = [tool.name for tool in tools] verbose_logger.info( "MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names ) - return result.tools + return list(tools) except asyncio.CancelledError: verbose_logger.warning("MCP client list_tools was cancelled") raise @@ -734,17 +740,17 @@ class MCPClient: """List available prompts from the server.""" verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_prompts_operation(session: ClientSession): - return await session.list_prompts() + async def _list_prompts_operation(session: ClientSession) -> tuple[Prompt, ...]: + return await list_all_prompts(session, self.server_url or "stdio") try: - result: Final = await self.run_with_session(_list_prompts_operation) - prompt_count: Final = len(result.prompts) - prompt_names: Final = [prompt.name for prompt in result.prompts] + prompts: Final = await self.run_with_session(_list_prompts_operation) + prompt_count: Final = len(prompts) + prompt_names: Final = [prompt.name for prompt in prompts] verbose_logger.info( - "MCP client listed %s tools from %s: %s", prompt_count, self.server_url or "stdio", prompt_names + "MCP client listed %s prompts from %s: %s", prompt_count, self.server_url or "stdio", prompt_names ) - return result.prompts + return list(prompts) except asyncio.CancelledError: verbose_logger.warning("MCP client list_prompts was cancelled") raise @@ -811,17 +817,17 @@ class MCPClient: """List available resources from the server.""" verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio") - async def _list_resources_operation(session: ClientSession): - return await session.list_resources() + async def _list_resources_operation(session: ClientSession) -> tuple[Resource, ...]: + return await list_all_resources(session, self.server_url or "stdio") try: - result: Final = await self.run_with_session(_list_resources_operation) - resource_count: Final = len(result.resources) - resource_names: Final = [resource.name for resource in result.resources] + resources: Final = await self.run_with_session(_list_resources_operation) + resource_count: Final = len(resources) + resource_names: Final = [resource.name for resource in resources] verbose_logger.info( "MCP client listed %s resources from %s: %s", resource_count, self.server_url or "stdio", resource_names ) - return result.resources + return list(resources) except asyncio.CancelledError: verbose_logger.warning("MCP client list_resources was cancelled") raise @@ -847,20 +853,20 @@ class MCPClient: """List available resource templates from the server.""" verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio") - async def _list_resource_templates_operation(session: ClientSession): - return await session.list_resource_templates() + async def _list_resource_templates_operation(session: ClientSession) -> tuple[ResourceTemplate, ...]: + return await list_all_resource_templates(session, self.server_url or "stdio") try: - result: Final = await self.run_with_session(_list_resource_templates_operation) - resource_template_count: Final = len(result.resourceTemplates) - resource_template_names: Final = [resourceTemplate.name for resourceTemplate in result.resourceTemplates] + resource_templates: Final = await self.run_with_session(_list_resource_templates_operation) + resource_template_count: Final = len(resource_templates) + resource_template_names: Final = [resource_template.name for resource_template in resource_templates] verbose_logger.info( "MCP client listed %s resource templates from %s: %s", resource_template_count, self.server_url or "stdio", resource_template_names, ) - return result.resourceTemplates + return list(resource_templates) except asyncio.CancelledError: verbose_logger.warning("MCP client list_resource_templates was cancelled") raise diff --git a/litellm/experimental_mcp_client/pagination.py b/litellm/experimental_mcp_client/pagination.py new file mode 100644 index 00000000000..8852715aba5 --- /dev/null +++ b/litellm/experimental_mcp_client/pagination.py @@ -0,0 +1,92 @@ +""" +Follows ``nextCursor`` on the paginated MCP list operations so a multi-page catalog is read in full. +""" + +from collections.abc import Awaitable, Callable, Sequence +from typing import Final, TypeVar + +from mcp import ClientSession, Resource +from mcp.types import PaginatedRequestParams, PaginatedResult, Prompt, ResourceTemplate +from mcp.types import Tool as MCPTool + +from litellm._logging import verbose_logger +from litellm.constants import MCP_LIST_MAX_PAGES + +TPage = TypeVar("TPage", bound=PaginatedResult) +TItem = TypeVar("TItem") + + +async def collect_pages( + fetch_page: Callable[[PaginatedRequestParams | None], Awaitable[TPage]], + items_of: Callable[[TPage], Sequence[TItem]], + *, + method: str, + server: str, + cursor: str | None = None, + seen_cursors: frozenset[str] = frozenset(), +) -> tuple[TItem, ...]: + page: Final = await fetch_page(None if cursor is None else PaginatedRequestParams(cursor=cursor)) + items: Final = tuple(items_of(page)) + next_cursor: Final = page.nextCursor + pages_read: Final = len(seen_cursors) + 1 + if next_cursor is None: + return items + if next_cursor in seen_cursors: + verbose_logger.warning( + "MCP %s from %s repeated cursor %r; returning the %s page(s) read so far", + method, + server, + next_cursor, + pages_read, + ) + return items + if pages_read >= MCP_LIST_MAX_PAGES: + verbose_logger.warning( + "MCP %s from %s still paginating after %s pages (LITELLM_MCP_LIST_MAX_PAGES); returning what was read", + method, + server, + pages_read, + ) + return items + rest: Final = await collect_pages( + fetch_page, + items_of, + method=method, + server=server, + cursor=next_cursor, + seen_cursors=seen_cursors | frozenset((next_cursor,)), + ) + return items + rest + + +async def list_all_tools(session: ClientSession, server: str) -> tuple[MCPTool, ...]: + return await collect_pages( + lambda params: session.list_tools(params=params), lambda page: page.tools, method="tools/list", server=server + ) + + +async def list_all_prompts(session: ClientSession, server: str) -> tuple[Prompt, ...]: + return await collect_pages( + lambda params: session.list_prompts(params=params), + lambda page: page.prompts, + method="prompts/list", + server=server, + ) + + +async def list_all_resources(session: ClientSession, server: str) -> tuple[Resource, ...]: + return await collect_pages( + lambda params: session.list_resources(params=params), + lambda page: page.resources, + method="resources/list", + server=server, + ) + + +async def list_all_resource_templates(session: ClientSession, server: str) -> tuple[ResourceTemplate, ...]: + return await collect_pages( + lambda params: session.list_resource_templates(params=params), + lambda page: page.resourceTemplates, + method="resources/templates/list", + server=server, + ) diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 30d50e2a74b..adaca0888aa 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -9,6 +9,7 @@ from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition +from litellm.experimental_mcp_client.pagination import list_all_tools from litellm.types.llms.anthropic import AnthropicMessagesTool from litellm.types.utils import ChatCompletionMessageToolCall @@ -103,10 +104,10 @@ async def load_mcp_tools( If format is set to "openai", the tools are converted to OpenAI API compatible tools. """ - tools: Final = await session.list_tools() + tools: Final = await list_all_tools(session, "upstream") if format == "openai": - return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools] - return tools.tools + return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools] + return list(tools) ######################################################## diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3efb6429326..3ca6b6c5f90 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1402,11 +1402,7 @@ if MCP_AVAILABLE: oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) async def _list_tools_operation(client): - async def _list_tools_session_operation(session): - return await session.list_tools() - - list_tools_response: Final = await client.run_with_session(_list_tools_session_operation) - list_tools_result: Final[list[MCPTool]] = list_tools_response.tools + list_tools_result: Final[list[MCPTool]] = await client.list_tools(raise_on_error=True) model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index fd7ab3afdab..3f501d3859a 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -20,8 +20,10 @@ from mcp.types import ( JSONRPCError, JSONRPCMessage, JSONRPCResponse, + ListToolsResult, ServerCapabilities, ) +from mcp.types import Tool as MCPTool # Add the parent directory to the path so we can import litellm @@ -740,8 +742,14 @@ class _ScriptedUpstream: error, the shape an upstream application uses to report its own failure. """ - def __init__(self, tools_list_error: ErrorData | None = None): + def __init__( + self, + tools_list_error: ErrorData | None = None, + tool_pages: tuple[tuple[MCPTool, ...], ...] = (), + ): self._tools_list_error = tools_list_error + self._tool_pages = tool_pages + self.tools_list_cursors: list[str | None] = [] self._to_client_tx, self._to_client_rx = anyio.create_memory_object_stream(10) self._from_client_tx, self._from_client_rx = anyio.create_memory_object_stream(10) self._task_group = None @@ -778,15 +786,37 @@ class _ScriptedUpstream: ) elif method == "tools/list" and self._tools_list_error is not None: await self._send(JSONRPCError(jsonrpc="2.0", id=request.id, error=self._tools_list_error)) + elif method == "tools/list" and self._tool_pages: + cursor = (request.params or {}).get("cursor") + self.tools_list_cursors.append(cursor) + page_index = int(cursor) if cursor else 0 + has_more = page_index + 1 < len(self._tool_pages) + page = ListToolsResult( + tools=list(self._tool_pages[page_index]), + nextCursor=str(page_index + 1) if has_more else None, + ) + await self._send( + JSONRPCResponse( + jsonrpc="2.0", + id=request.id, + result=page.model_dump(by_alias=True, mode="json", exclude_none=True), + ) + ) class _ScriptedClient(MCPClient): """An MCPClient whose transport is a scripted in-memory upstream instead of a real connection, so the real ``ClientSession`` and its real timeout machinery are what run.""" - def __init__(self, *, timeout: float, tools_list_error: ErrorData | None = None): + def __init__( + self, + *, + timeout: float, + tools_list_error: ErrorData | None = None, + tool_pages: tuple[tuple[MCPTool, ...], ...] = (), + ): super().__init__(server_url="http://upstream.local/mcp", timeout=timeout) - self._upstream = _ScriptedUpstream(tools_list_error=tools_list_error) + self._upstream = _ScriptedUpstream(tools_list_error=tools_list_error, tool_pages=tool_pages) def _create_transport_context(self): return self._upstream, None @@ -821,6 +851,22 @@ async def test_list_tools_fails_on_its_own_timeout_when_the_upstream_never_answe assert list_fault_http_status(fault) == 504 +@pytest.mark.asyncio +async def test_list_tools_follows_tools_list_pagination_across_the_whole_catalog(): + """An upstream that pages tools/list (72 tools, 30 per page) must have every page read within the + one session, each request carrying the cursor the previous page returned. Reading only the first + page made 42 tools invisible to the proxy and every call to them fail as unknown.""" + tools = tuple( + MCPTool(name=f"tool_{i:02d}", inputSchema={"type": "object", "properties": {}}) for i in range(72) + ) + client = _ScriptedClient(timeout=30, tool_pages=(tools[:30], tools[30:60], tools[60:])) + + listed = await asyncio.wait_for(client.list_tools(raise_on_error=True), timeout=10) + + assert [tool.name for tool in listed] == [tool.name for tool in tools] + assert client._upstream.tools_list_cursors == [None, "1", "2"] + + @pytest.mark.asyncio async def test_upstream_json_rpc_error_408_is_not_reported_as_a_client_timeout(): """The SDK reports its own elapsed read timeout and relays an upstream JSON-RPC error through diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py new file mode 100644 index 00000000000..f588fdd1eee --- /dev/null +++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py @@ -0,0 +1,80 @@ +import logging + +import pytest +from mcp.types import ListToolsResult, PaginatedRequestParams +from mcp.types import Tool as MCPTool + +import litellm.experimental_mcp_client.pagination as pagination_module +from litellm.experimental_mcp_client.pagination import collect_pages + + +def _tool(index: int) -> MCPTool: + return MCPTool(name=f"tool_{index:02d}", inputSchema={"type": "object", "properties": {}}) + + +class _PagedTools: + """A tools/list upstream serving ``total`` tools ``page_size`` at a time, cursors being offsets.""" + + def __init__(self, total: int, page_size: int): + self._tools = tuple(_tool(i) for i in range(total)) + self._page_size = page_size + self.cursors_seen: list[str | None] = [] + + async def fetch(self, params: PaginatedRequestParams | None) -> ListToolsResult: + cursor = params.cursor if params is not None else None + self.cursors_seen.append(cursor) + start = int(cursor) if cursor else 0 + end = start + self._page_size + return ListToolsResult( + tools=list(self._tools[start:end]), + nextCursor=str(end) if end < len(self._tools) else None, + ) + + +@pytest.mark.asyncio +async def test_collect_pages_follows_next_cursor_until_exhausted(): + upstream = _PagedTools(total=72, page_size=30) + + tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s") + + assert [t.name for t in tools] == [f"tool_{i:02d}" for i in range(72)] + assert upstream.cursors_seen == [None, "30", "60"], "each page must be requested with the cursor the previous one returned" + + +@pytest.mark.asyncio +async def test_collect_pages_single_page_makes_one_request(): + upstream = _PagedTools(total=5, page_size=30) + + tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s") + + assert len(tools) == 5 + assert upstream.cursors_seen == [None] + + +@pytest.mark.asyncio +async def test_collect_pages_stops_on_a_repeated_cursor_and_keeps_what_it_read(caplog): + calls: list[str | None] = [] + + async def fetch(params: PaginatedRequestParams | None) -> ListToolsResult: + calls.append(params.cursor if params else None) + return ListToolsResult(tools=[_tool(len(calls))], nextCursor="same") + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + tools = await collect_pages(fetch, lambda page: page.tools, method="tools/list", server="s") + + assert calls == [None, "same"], "the cursor must be followed once and refused the second time it comes back" + assert len(tools) == 2 + assert any("repeated cursor" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_collect_pages_honors_the_page_cap(monkeypatch, caplog): + monkeypatch.setattr(pagination_module, "MCP_LIST_MAX_PAGES", 3) + upstream = _PagedTools(total=1000, page_size=10) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s") + + assert len(upstream.cursors_seen) == 3 + assert len(tools) == 30 + assert any("LITELLM_MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records) From 26b48d58919e5021b9b251339bf5c720a3e3649e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 21:00:27 +0000 Subject: [PATCH 005/246] refactor(mcp): keep list pagination within type-discipline budget and ratchet LIT001 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 56 ++++++++++--------- .../mcp_server/rest_endpoints.py | 4 +- .../test_pagination.py | 4 +- type-discipline-budget.json | 2 +- 4 files changed, 35 insertions(+), 31 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 814074d35a4..46d8823e439 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -377,29 +377,31 @@ class MCPClient: return provided_env # Minimal allowlist of safe/standard environment variables - safe_keys: Final = { - "PATH", - "HOME", - "USER", - "LOGNAME", - "TMPDIR", - "TMP", - "TEMP", - "SHELL", - "LANG", - "LC_ALL", - # Node/Package manager caches - "NPM_CONFIG_CACHE", - "PNPM_HOME", - "XDG_CACHE_HOME", - "XDG_CONFIG_HOME", - "XDG_DATA_HOME", - # System info - "SYSTEMROOT", - "COMSPEC", - "PATHEXT", - "WINDIR", - } + safe_keys: Final = frozenset( + { + "PATH", + "HOME", + "USER", + "LOGNAME", + "TMPDIR", + "TMP", + "TEMP", + "SHELL", + "LANG", + "LC_ALL", + # Node/Package manager caches + "NPM_CONFIG_CACHE", + "PNPM_HOME", + "XDG_CACHE_HOME", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + # System info + "SYSTEMROOT", + "COMSPEC", + "PATHEXT", + "WINDIR", + } + ) safe_env: Final = {} for key in safe_keys: @@ -615,7 +617,7 @@ class MCPClient: try: tools: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) tool_count: Final = len(tools) - tool_names: Final = [tool.name for tool in tools] + tool_names: Final = tuple(tool.name for tool in tools) verbose_logger.info( "MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names ) @@ -746,7 +748,7 @@ class MCPClient: try: prompts: Final = await self.run_with_session(_list_prompts_operation) prompt_count: Final = len(prompts) - prompt_names: Final = [prompt.name for prompt in prompts] + prompt_names: Final = tuple(prompt.name for prompt in prompts) verbose_logger.info( "MCP client listed %s prompts from %s: %s", prompt_count, self.server_url or "stdio", prompt_names ) @@ -823,7 +825,7 @@ class MCPClient: try: resources: Final = await self.run_with_session(_list_resources_operation) resource_count: Final = len(resources) - resource_names: Final = [resource.name for resource in resources] + resource_names: Final = tuple(resource.name for resource in resources) verbose_logger.info( "MCP client listed %s resources from %s: %s", resource_count, self.server_url or "stdio", resource_names ) @@ -859,7 +861,7 @@ class MCPClient: try: resource_templates: Final = await self.run_with_session(_list_resource_templates_operation) resource_template_count: Final = len(resource_templates) - resource_template_names: Final = [resource_template.name for resource_template in resource_templates] + resource_template_names: Final = tuple(resource_template.name for resource_template in resource_templates) verbose_logger.info( "MCP client listed %s resource templates from %s: %s", resource_template_count, diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3ca6b6c5f90..beee7c1db78 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,6 +1,6 @@ import asyncio import importlib -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -1402,7 +1402,7 @@ if MCP_AVAILABLE: oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) async def _list_tools_operation(client): - list_tools_result: Final[list[MCPTool]] = await client.list_tools(raise_on_error=True) + list_tools_result: Final[Sequence[MCPTool]] = await client.list_tools(raise_on_error=True) model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py index f588fdd1eee..a76f410ac3c 100644 --- a/tests/test_litellm/experimental_mcp_client/test_pagination.py +++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py @@ -38,7 +38,9 @@ async def test_collect_pages_follows_next_cursor_until_exhausted(): tools = await collect_pages(upstream.fetch, lambda page: page.tools, method="tools/list", server="s") assert [t.name for t in tools] == [f"tool_{i:02d}" for i in range(72)] - assert upstream.cursors_seen == [None, "30", "60"], "each page must be requested with the cursor the previous one returned" + assert upstream.cursors_seen == [None, "30", "60"], ( + "each page must be requested with the cursor the previous one returned" + ) @pytest.mark.asyncio diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3d2e97d55a5..b0c3cc7f9fd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22367 + "limit": 22366 }, "LIT002": { "limit": 26777 From d32f8a07c88ed165e55ce944026504a1cd15a327 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 21:15:10 +0000 Subject: [PATCH 006/246] fix(mcp): make list page cap a plain constant and use a real ListToolsResult in the unit mock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 +- litellm/experimental_mcp_client/pagination.py | 2 +- tests/mcp_tests/test_mcp_client_unit.py | 6 ++---- .../test_litellm/experimental_mcp_client/test_pagination.py | 2 +- 4 files changed, 5 insertions(+), 7 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 11f35177636..07914934495 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -136,7 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0" MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) -MCP_LIST_MAX_PAGES: Final = int(os.getenv("LITELLM_MCP_LIST_MAX_PAGES", "100")) +MCP_LIST_MAX_PAGES: Final = 100 # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. diff --git a/litellm/experimental_mcp_client/pagination.py b/litellm/experimental_mcp_client/pagination.py index 8852715aba5..85f46268f92 100644 --- a/litellm/experimental_mcp_client/pagination.py +++ b/litellm/experimental_mcp_client/pagination.py @@ -42,7 +42,7 @@ async def collect_pages( return items if pages_read >= MCP_LIST_MAX_PAGES: verbose_logger.warning( - "MCP %s from %s still paginating after %s pages (LITELLM_MCP_LIST_MAX_PAGES); returning what was read", + "MCP %s from %s still paginating after %s pages (MCP_LIST_MAX_PAGES); returning what was read", method, server, pages_read, diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index aadaadd510e..ef4231fe1d9 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/tests/mcp_tests/test_mcp_client_unit.py @@ -11,7 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch, ANY import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient from litellm.types.mcp import MCPAuth, MCPTransport -from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult +from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult, ListToolsResult def test_mcp_client_uses_configurable_default_timeout(): @@ -174,9 +174,7 @@ class TestMCPClientUnitTests: }, ) ] - mock_result = MagicMock() - mock_result.tools = mock_tools - mock_session_instance.list_tools.return_value = mock_result + mock_session_instance.list_tools.return_value = ListToolsResult(tools=mock_tools) client = MCPClient("http://example.com") result = await client.list_tools() diff --git a/tests/test_litellm/experimental_mcp_client/test_pagination.py b/tests/test_litellm/experimental_mcp_client/test_pagination.py index a76f410ac3c..93952c176a8 100644 --- a/tests/test_litellm/experimental_mcp_client/test_pagination.py +++ b/tests/test_litellm/experimental_mcp_client/test_pagination.py @@ -79,4 +79,4 @@ async def test_collect_pages_honors_the_page_cap(monkeypatch, caplog): assert len(upstream.cursors_seen) == 3 assert len(tools) == 30 - assert any("LITELLM_MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records) + assert any("MCP_LIST_MAX_PAGES" in record.getMessage() for record in caplog.records) From 391b8a265b7c5285b2640361268ec67b72c143f3 Mon Sep 17 00:00:00 2001 From: Atharva-Kanherkar <142440039+Atharva-Kanherkar@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:40:43 +0530 Subject: [PATCH 007/246] fix(responses): stream one lifecycle across MCP auto-execute rounds Each auto-executed MCP round is a distinct upstream response, but the client reads one stream. The follow-up round's response.created, response.in_progress, and the interim response.completed were forwarded as-is, so one SSE body carried two lifecycles and output_index restarted at zero, which aborts accumulating clients such as the OpenAI SDK's responses.stream() before the final answer arrives. Fold the rounds into one public lifecycle: drop the openers of follow-up rounds, hold back the completed event of a round whose tool calls the gateway executes, shift later output indexes past the items already emitted, and list every round's items on the single final response.completed. Each mcp_call item is announced with output_item.added, keeps one item id across its events, and owns its own output_index. Sequence numbers stay strictly increasing when a round or a gateway event restarts numbering. The final round's response id is kept on response.completed so previous_response_id continuation still works. --- .../responses/mcp/mcp_streaming_iterator.py | 212 ++++++++++++++---- 1 file changed, 166 insertions(+), 46 deletions(-) diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index ca12b3e7cc3..ac158a32371 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -34,6 +34,16 @@ else: MAX_MCP_TOOL_CALL_ROUNDS: Final = 5 +def _output_items(response: ResponsesAPIResponse) -> Sequence[object]: + """Read a response's output items as plain objects; the field is a wide union of item models.""" + return tuple(cast("Sequence[object]", response.output)) # cast-ok: items are only carried, never inspected + + +def _set_event_field(event: ResponsesAPIStreamingResponse, name: str, value: object) -> None: + """Events are pydantic models with extra fields allowed, so any event type can carry the field.""" + setattr(event, name, value) + + async def create_mcp_list_tools_events( mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], user_api_key_auth: "UserAPIKeyAuth | None", @@ -170,6 +180,7 @@ def create_mcp_call_events( result: str | None = None, base_item_id: str | None = None, sequence_start: int = 1, + output_index: int = 0, ) -> list[ResponsesAPIStreamingResponse]: """Create MCP call events following OpenAI's specification""" events: Final[list[ResponsesAPIStreamingResponse]] = [] @@ -179,7 +190,7 @@ def create_mcp_call_events( in_progress_event: Final = MCPCallInProgressEvent( type=ResponsesAPIStreamEvents.MCP_CALL_IN_PROGRESS, sequence_number=sequence_start, - output_index=0, + output_index=output_index, item_id=item_id, ) events.append(in_progress_event) @@ -187,7 +198,7 @@ def create_mcp_call_events( # MCP call arguments delta event (streaming the arguments) arguments_delta_event: Final = MCPCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, - output_index=0, + output_index=output_index, item_id=item_id, delta=arguments, # JSON string with arguments sequence_number=sequence_start + 1, @@ -197,7 +208,7 @@ def create_mcp_call_events( # MCP call arguments done event arguments_done_event: Final = MCPCallArgumentsDoneEvent( type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DONE, - output_index=0, + output_index=output_index, item_id=item_id, arguments=arguments, # Complete JSON string with finalized arguments sequence_number=sequence_start + 2, @@ -210,7 +221,7 @@ def create_mcp_call_events( type=ResponsesAPIStreamEvents.MCP_CALL_COMPLETED, sequence_number=sequence_start + 3, item_id=item_id, - output_index=0, + output_index=output_index, ) events.append(completed_event) @@ -219,7 +230,7 @@ def create_mcp_call_events( output_item_done_event: Final = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, - output_index=0, + output_index=output_index, item=BaseLiteLLMOpenAIResponseObject( **{ "id": item_id, @@ -239,7 +250,7 @@ def create_mcp_call_events( type=ResponsesAPIStreamEvents.MCP_CALL_FAILED, sequence_number=sequence_start + 3, item_id=item_id, - output_index=0, + output_index=output_index, ) events.append(failed_event) @@ -330,6 +341,17 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._error_event_emitted = False self._last_sequence_number = 0 + # Every auto-execute round is a distinct upstream response, but the + # client is reading one stream. Fold the rounds into one public + # lifecycle: one response.created, one response.completed whose + # output holds every round's items, and output indexes that are + # never reused for a different item. + self._round_index = 0 + self._output_index_offset = 0 + self._round_max_output_index = -1 + self._composed_output: list[object] = [] # mutable-ok: grows as each round finishes + self._pending_mcp_call_items: list[dict[str, object]] = [] # mutable-ok: grows per executed tool + def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" @@ -415,8 +437,14 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): async def __anext__(self) -> ResponsesAPIStreamingResponse: chunk: Final = await self._anext_impl() sequence_number: Final = getattr(chunk, "sequence_number", None) - if isinstance(sequence_number, int) and sequence_number > self._last_sequence_number: - self._last_sequence_number = sequence_number + if isinstance(sequence_number, int): + # Follow-up rounds and gateway events restart their numbering. + # Keep the public stream strictly increasing. + if sequence_number <= self._last_sequence_number and self._last_sequence_number > 0: + self._last_sequence_number += 1 + _set_event_field(chunk, "sequence_number", self._last_sequence_number) + else: + self._last_sequence_number = max(self._last_sequence_number, sequence_number) return chunk async def _anext_impl(self) -> ResponsesAPIStreamingResponse: @@ -472,7 +500,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): await self._create_follow_up_iterator() if self.base_iterator is not None: self.phase = "continue_initial_response" - return await self.__anext__() + return await self._anext_impl() self.phase = "finished" if self._stream_error is not None and not self._error_event_emitted: self._error_event_emitted = True @@ -530,17 +558,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if chunk_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED: self.initial_events_emitted = True self.phase = "mcp_discovery" - return chunk + return await self._compose_round_chunk(chunk) - # If auto-execution is enabled, check for completed responses - if self.should_auto_execute and self._is_response_completed(chunk): - response_obj = getattr(chunk, "response", None) - if isinstance(response_obj, ResponsesAPIResponse): - self.collected_response = response_obj - self.phase = "tool_execution" - await self._generate_tool_execution_events() - - return chunk + # None means the chunk was folded into the single public + # lifecycle; fall through so phase 4 runs the follow-up. + return await self._compose_round_chunk(chunk) except StopAsyncIteration: if self.should_auto_execute and self.collected_response: self.phase = "tool_execution" @@ -566,6 +588,69 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): chunk_type: Final[object] = getattr(chunk, "type", None) return chunk_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + def _follow_up_pending(self) -> bool: + """True when the current round's tool calls were executed and a follow-up round will run.""" + return self.collected_response is not None and self.collected_response is self._tool_results_for_response + + def _round_output_width(self, response: ResponsesAPIResponse) -> int: + """How many output indexes this round used, counting items it streamed but never listed.""" + return max(len(_output_items(response)), self._round_max_output_index + 1) + + def _absorb_round(self, response: ResponsesAPIResponse) -> None: + """Bank a finished round's items so the final response.completed can list them.""" + width: Final = self._round_output_width(response) + self._composed_output.extend(_output_items(response)) + self._composed_output.extend(self._pending_mcp_call_items) + self._output_index_offset += width + len(self._pending_mcp_call_items) + self._pending_mcp_call_items = [] + self._round_max_output_index = -1 + + async def _compose_round_chunk(self, chunk: ResponsesAPIStreamingResponse) -> ResponsesAPIStreamingResponse | None: + """ + Fold one round's event into the single public lifecycle. + + Returns None when the event must not reach the client: the lifecycle + openers of a follow-up round, and the response.completed of a round + whose tool calls the gateway executes itself. Shifts output_index on + follow-up rounds past the items already emitted, and lists every + round's items on the final response.completed. + """ + chunk_type: Final[object] = getattr(chunk, "type", None) + if self._round_index > 0 and chunk_type in ( + ResponsesAPIStreamEvents.RESPONSE_CREATED, + ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, + ): + return None + + output_index: Final[object] = getattr(chunk, "output_index", None) + if isinstance(output_index, int): + self._round_max_output_index = max(self._round_max_output_index, output_index) + if self._output_index_offset: + _set_event_field(chunk, "output_index", output_index + self._output_index_offset) + + if not (self.should_auto_execute and self._is_response_completed(chunk)): + return chunk + + response_obj: Final[object] = getattr(chunk, "response", None) + if isinstance(response_obj, ResponsesAPIResponse): + self.collected_response = response_obj + # Move to tool execution phase after this chunk + self.phase = "tool_execution" + await self._generate_tool_execution_events() + + if not isinstance(response_obj, ResponsesAPIResponse): + return chunk + if self._follow_up_pending(): + self._absorb_round(response_obj) + return None + if self._composed_output: + merged_output: Final[list[object]] = [ # mutable-ok: the response model declares output as a list + *self._composed_output, + *_output_items(response_obj), + ] + _set_event_field(chunk, "response", response_obj.model_copy(update={"output": merged_output})) + return chunk + async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: """ Process a chunk from the base iterator with response ID consistency enforcement. @@ -593,17 +678,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): ) response_obj.id = self._cached_response_id - # If auto-execution is enabled, check for completed responses - if self.should_auto_execute and self._is_response_completed(chunk): - # Collect the response for tool execution - response_obj = getattr(chunk, "response", None) - if isinstance(response_obj, ResponsesAPIResponse): - self.collected_response = response_obj - # Move to tool execution phase after emitting this chunk - self.phase = "tool_execution" - await self._generate_tool_execution_events() - - return chunk + composed: Final = await self._compose_round_chunk(chunk) + if composed is None: + # The chunk stays internal; hand the next public event back instead. + return await self._anext_impl() + return composed async def _create_initial_response_iterator(self) -> None: """Create the initial response iterator by making the first LLM call""" @@ -667,6 +746,16 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return self.tool_call_round += 1 + # Each executed tool is one mcp_call output item of the single + # public response. Announce it at an output_index past the items + # this round already streamed, and keep that item id for the + # completion events below. + from litellm.types.llms.openai import OutputItemAddedEvent + + next_output_index = self._output_index_offset + self._round_output_width( # rebind-ok: advances per item + self.collected_response + ) + call_items: Final[dict[str, tuple[str, int]]] = {} # mutable-ok: filled per tool call as events queue for tool_call in tool_calls: ( tool_name, @@ -674,14 +763,36 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_call_id, ) = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) if tool_name and tool_call_id: + item_id = f"mcp_{uuid.uuid4().hex[:8]}" + output_index = next_output_index + next_output_index += 1 + call_items[tool_call_id] = (item_id, output_index) + self.tool_execution_events.append( + OutputItemAddedEvent.model_validate( + { + "type": ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + "sequence_number": len(self.tool_execution_events) + 1, + "output_index": output_index, + "item": { + "id": item_id, + "type": "mcp_call", + "status": "in_progress", + "arguments": tool_arguments or "{}", + "name": tool_name, + "server_label": "litellm", + }, + } + ) + ) # Create MCP call events for this tool execution call_events = create_mcp_call_events( tool_name=tool_name, tool_call_id=tool_call_id, arguments=tool_arguments or "{}", # JSON string with arguments result=None, # Will be set after execution - base_item_id=f"mcp_{uuid.uuid4().hex[:8]}", + base_item_id=item_id, sequence_start=len(self.tool_execution_events) + 1, + output_index=output_index, ) # Add the in_progress and arguments events (not the completed event yet) self.tool_execution_events.extend(call_events[:-1]) @@ -719,37 +830,45 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_arguments = args or "{}" break - item_id = f"mcp_{uuid.uuid4().hex[:8]}" + if tool_call_id in call_items: + item_id, output_index = call_items[tool_call_id] + else: + item_id = f"mcp_{uuid.uuid4().hex[:8]}" + output_index = next_output_index + next_output_index += 1 # Create the completion event completed_event = MCPCallCompletedEvent( type=ResponsesAPIStreamEvents.MCP_CALL_COMPLETED, sequence_number=len(self.tool_execution_events) + 1, item_id=item_id, - output_index=0, + output_index=output_index, ) self.tool_execution_events.append(completed_event) # Create output_item.done event with the tool call result from litellm.types.llms.openai import OutputItemDoneEvent + mcp_call_item = BaseLiteLLMOpenAIResponseObject( + **{ + "id": item_id, + "type": "mcp_call", + "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", + "arguments": tool_arguments, + "error": None, + "name": tool_name, + "output": result_text, + "server_label": "litellm", # or extract from tool config + } + ) output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, - output_index=0, - item=BaseLiteLLMOpenAIResponseObject( - **{ - "id": item_id, - "type": "mcp_call", - "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", - "arguments": tool_arguments, - "error": None, - "name": tool_name, - "output": result_text, - "server_label": "litellm", # or extract from tool config - } - ), + output_index=output_index, + item=mcp_call_item, ) self.tool_execution_events.append(output_item_done_event) + # The response model accepts output items as dicts, not as the generic event object. + self._pending_mcp_call_items.append(mcp_call_item.model_dump()) # Store tool results for follow-up call self.tool_results = tool_results @@ -824,6 +943,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.base_iterator = follow_up_response self.collected_response = None self._cached_response_id = None + self._round_index += 1 except Exception as e: verbose_logger.error("Error creating follow-up iterator: %s", e) From 13ce30488ff6eed68c213a80a41a38979a279e59 Mon Sep 17 00:00:00 2001 From: Atharva-Kanherkar <142440039+Atharva-Kanherkar@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:40:43 +0530 Subject: [PATCH 008/246] test(responses): cover the single public lifecycle for MCP rounds Assert one response.created and one response.completed across rounds, dense output indexes for the function call, the mcp_call item, and the final message, stable mcp_call item ids, strictly increasing sequence numbers, a serializable merged output, and that a stream without auto-execution is forwarded unchanged. Update the two existing tests that counted one completed event per internal round. --- .../mcp/test_mcp_streaming_iterator.py | 128 +++++++++++++++++- 1 file changed, 122 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py index 5001589ce54..57ccbbfa493 100644 --- a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py @@ -57,6 +57,10 @@ def _text_message(text: str): return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} +def _item_type(item) -> str: + return item["type"] if isinstance(item, dict) else item.type + + def _tool_call_stream(call_id: str, tool_name: str, response_id: str = "resp-1") -> _FakeAsyncStream: return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)], response_id=response_id)]) @@ -134,11 +138,19 @@ async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeyp assert iterator.tool_call_round == 2 # The stream reached round 3 and produced the final text response instead - # of stopping after round 1 or round 2. + # of stopping after round 1 or round 2. The client sees one lifecycle whose + # final output lists every round's items in order. completed_chunks = [c for c in chunks if getattr(c, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED] - assert len(completed_chunks) == 3 + assert len(completed_chunks) == 1 final_output = completed_chunks[-1].response.output - assert final_output[0]["content"][0]["text"] == "Here's what I found after retrying." + assert [_item_type(item) for item in final_output] == [ + "function_call", + "mcp_call", + "function_call", + "mcp_call", + "message", + ] + assert final_output[-1]["content"][0]["text"] == "Here's what I found after retrying." @pytest.mark.asyncio @@ -207,7 +219,7 @@ async def test_continuation_id_is_final_round_not_interim_tool_call(monkeypatch) chunks = [chunk async for chunk in iterator] completed = [c for c in chunks if getattr(c, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED] - assert completed[-1].response.output[0]["content"][0]["text"] == "The first item is Alpha." + assert completed[-1].response.output[-1]["content"][0]["text"] == "The first item is Alpha." assert completed[-1].response.id == "resp-final" assert completed[-1].response.id != "resp-interim" @@ -280,7 +292,9 @@ async def test_streaming_follow_up_replays_reasoning_when_store_is_false(monkeyp base_iterator=_FakeAsyncStream( [ _output_item_added_chunk(), - _completed_chunk([_reasoning_item("gAAAAA-opaque-blob"), _function_call("call_1", "read_wiki_contents")]), + _completed_chunk( + [_reasoning_item("gAAAAA-opaque-blob"), _function_call("call_1", "read_wiki_contents")] + ), ] ), mcp_events=[], @@ -316,7 +330,9 @@ async def test_streaming_follow_up_keeps_previous_response_id_when_stored(monkey base_iterator=_FakeAsyncStream( [ _output_item_added_chunk(), - _completed_chunk([_reasoning_item("gAAAAA-opaque-blob"), _function_call("call_1", "read_wiki_contents")]), + _completed_chunk( + [_reasoning_item("gAAAAA-opaque-blob"), _function_call("call_1", "read_wiki_contents")] + ), ] ), mcp_events=[], @@ -336,3 +352,103 @@ async def test_streaming_follow_up_keeps_previous_response_id_when_stored(monkey follow_up_kwargs = aresponses_mock.call_args_list[0].kwargs assert follow_up_kwargs["previous_response_id"] == "resp_prev" assert not [item for item in follow_up_kwargs["input"] if item.get("type") == "reasoning"] + + +def _event(event_type, **fields): + return SimpleNamespace(type=event_type, **fields) + + +def _lifecycle_round(response_id: str, item: dict, sequence_start: int = 0): + """One upstream Responses round as a provider streams it: its own id, indexes from 0, numbering from 0.""" + return [ + _event( + ResponsesAPIStreamEvents.RESPONSE_CREATED, + response=ResponsesAPIResponse(id=response_id, created_at=0, output=[]), + sequence_number=sequence_start, + ), + _event( + ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=0, item=item, sequence_number=sequence_start + 1 + ), + _event( + ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, item=item, sequence_number=sequence_start + 2 + ), + _completed_chunk([item], response_id=response_id), + ] + + +@pytest.mark.asyncio +async def test_auto_execute_rounds_share_one_public_lifecycle(monkeypatch): + """ + Every auto-execute round is a distinct upstream response, but the client + reads one stream. It must see one response.created, one response.completed, + and no output_index reused for a different item, otherwise accumulating + clients such as the OpenAI SDK's responses.stream() abort mid-stream. + """ + _mock_mcp_environment(monkeypatch) + + follow_up = _FakeAsyncStream(_lifecycle_round("resp-final", _text_message("Alpha."))) + monkeypatch.setattr(responses_main_module, "aresponses", AsyncMock(side_effect=[follow_up])) + + iterator = _make_iterator(_lifecycle_round("resp-interim", _function_call("call_1", "read_wiki_contents"))) + chunks = [chunk async for chunk in iterator] + types = [chunk.type for chunk in chunks] + + assert types.count(ResponsesAPIStreamEvents.RESPONSE_CREATED) == 1 + assert types.count(ResponsesAPIStreamEvents.RESPONSE_COMPLETED) == 1 + assert types[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + + # The function call, the gateway's mcp_call, and the final message each own an index. + added = [c for c in chunks if c.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] + assert [(c.output_index, _item_type(c.item)) for c in added] == [ + (0, "function_call"), + (1, "mcp_call"), + (2, "message"), + ] + mcp_item_ids = {c.item_id for c in chunks if c.type == ResponsesAPIStreamEvents.MCP_CALL_IN_PROGRESS} + mcp_done = [ + c for c in chunks if c.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and _item_type(c.item) == "mcp_call" + ] + assert [c.output_index for c in mcp_done] == [1] + assert {c.item.id for c in mcp_done} == mcp_item_ids + round_two = [ + c for c in chunks if getattr(c, "item_id", None) is None and getattr(c, "output_index", None) is not None + ] + assert max(c.output_index for c in round_two) == 2 + + # The single completed event lists every round's items and keeps the final round's id for continuation. + completed = chunks[-1] + assert completed.response.id == "resp-final" + assert [_item_type(item) for item in completed.response.output] == ["function_call", "mcp_call", "message"] + assert completed.response.output[-1]["content"][0]["text"] == "Alpha." + # The proxy serializes every chunk; the merged output must still be a valid response. + assert '"type":"mcp_call"' in completed.response.model_dump_json(exclude_none=True, exclude_unset=True) + + # Numbering stays strictly increasing across rounds and gateway events. + sequence_numbers = [c.sequence_number for c in chunks if getattr(c, "sequence_number", None) is not None] + assert sequence_numbers == sorted(sequence_numbers) + assert len(set(sequence_numbers)) == len(sequence_numbers) + + +@pytest.mark.asyncio +async def test_stream_without_auto_execute_is_forwarded_unchanged(monkeypatch): + """With approval required there is one round, and it passes through untouched.""" + _mock_mcp_environment(monkeypatch) + aresponses_mock = AsyncMock() + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + upstream = _lifecycle_round("resp-1", _function_call("call_1", "read_wiki_contents")) + iterator = MCPEnhancedStreamingIterator( + base_iterator=_FakeAsyncStream(list(upstream)), + mcp_events=[], + tool_server_map={"read_wiki_contents": "deepwiki"}, + mcp_tools_with_litellm_proxy=[{"require_approval": "always"}], + user_api_key_auth=None, + original_request_params={"model": "gpt-4", "input": "hi", "tools": [{"type": "mcp"}]}, + ) + + chunks = [chunk async for chunk in iterator] + + assert chunks == upstream + assert [c.output_index for c in chunks if hasattr(c, "output_index")] == [0, 0] + assert [c.sequence_number for c in chunks if hasattr(c, "sequence_number")] == [0, 1, 2] + aresponses_mock.assert_not_called() From 146e5c492b1bee31a57c77a8fafff74122a13df9 Mon Sep 17 00:00:00 2001 From: Atharva-Kanherkar <142440039+Atharva-Kanherkar@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:49:04 +0530 Subject: [PATCH 009/246] fix(responses): keep the MCP lifecycle change within the type-discipline budget Mark the dict literals handed straight to model constructors, clear the pending mcp_call list in place, and bind the merged response before setting it on the event, so LIT002 stays at or below its base count. --- litellm/responses/mcp/mcp_streaming_iterator.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index ac158a32371..1c892eec803 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -602,7 +602,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._composed_output.extend(_output_items(response)) self._composed_output.extend(self._pending_mcp_call_items) self._output_index_offset += width + len(self._pending_mcp_call_items) - self._pending_mcp_call_items = [] + self._pending_mcp_call_items.clear() self._round_max_output_index = -1 async def _compose_round_chunk(self, chunk: ResponsesAPIStreamingResponse) -> ResponsesAPIStreamingResponse | None: @@ -648,7 +648,10 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): *self._composed_output, *_output_items(response_obj), ] - _set_event_field(chunk, "response", response_obj.model_copy(update={"output": merged_output})) + merged_response: Final = response_obj.model_copy( + update={"output": merged_output} # mutable-ok: pydantic's update argument must be a dict + ) + _set_event_field(chunk, "response", merged_response) return chunk async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: @@ -769,11 +772,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): call_items[tool_call_id] = (item_id, output_index) self.tool_execution_events.append( OutputItemAddedEvent.model_validate( - { + { # mutable-ok: consumed once by model_validate "type": ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "sequence_number": len(self.tool_execution_events) + 1, "output_index": output_index, - "item": { + "item": { # mutable-ok: consumed once by model_validate "id": item_id, "type": "mcp_call", "status": "in_progress", @@ -850,7 +853,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): from litellm.types.llms.openai import OutputItemDoneEvent mcp_call_item = BaseLiteLLMOpenAIResponseObject( - **{ + **{ # mutable-ok: consumed once by the model constructor "id": item_id, "type": "mcp_call", "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", From fa966ca2d0bc9bca963112c8befbda5997f306ef Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 08:10:30 +0000 Subject: [PATCH 010/246] refactor(types): replace Any with real types across 20 backend files Fifth round of basedpyright Any reduction. Every change is typing-only and leaves runtime behavior identical. Prisma table access now goes through the PrismaTableRepository and TableActions protocols the repo already has, instead of reading untyped attributes off prisma_client.db. Payload and parameter annotations move from dict[str, Any] to dict[str, object] or Mapping[str, object]. Guardrail constructors that took **kwargs: Any now take Unpack of a PEP 728 TypedDict, the same _CustomGuardrailOptions shape three other guardrails already use. Calls into the OpenAI and Azure assistants SDKs pass explicit keywords rather than splatting a dict, so the arguments are checked against the real SDK signatures. --- litellm/llms/azure/assistants.py | 9 +-- litellm/llms/ollama/chat/transformation.py | 2 +- .../llms/ollama/completion/transformation.py | 2 +- litellm/llms/openai/openai.py | 61 ++++++++++++++----- litellm/llms/vertex_ai/fine_tuning/handler.py | 5 +- .../vertex_imagen_transformation.py | 25 ++++---- .../vertex_gemini_transformation.py | 6 +- litellm/proxy/_experimental/mcp_server/db.py | 8 ++- .../proxy/common_utils/reset_budget_job.py | 15 +++-- .../cisco_ai_defense/cisco_ai_defense.py | 7 ++- .../guardrail_hooks/deepkeep/deepkeep.py | 14 +++-- .../guardrail_hooks/headroom/headroom.py | 13 ++-- .../model_armor/model_armor.py | 6 +- .../auto_router_endpoints.py | 10 +-- .../spend_tracking/ptu_flat_cost_rollup.py | 30 +++++++-- litellm/proxy/utils.py | 11 ++-- litellm/realtime_api/main.py | 2 +- litellm/repositories/model_repository.py | 20 ++---- litellm/repositories/team_repository.py | 42 ++++++++++++- .../router_utils/fallback_event_handlers.py | 11 ++-- 20 files changed, 199 insertions(+), 100 deletions(-) diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index f7b419405ac..a4742a25a87 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -1,5 +1,5 @@ from collections.abc import Coroutine, Iterable -from typing import Any, Final, Literal, TypedDict +from typing import Final, Literal, TypedDict import httpx from openai import AsyncAzureOpenAI, AzureOpenAI @@ -715,7 +715,8 @@ class AzureAssistantsAPI(BaseAzureLLM): event_handler: AssistantEventHandler | None, litellm_params: dict | None = None, ) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]: - data: Final[dict[str, Any]] = { + stream_fn: Final = client.beta.threads.runs.stream + base_data: Final[_RunThreadStreamData] = { "thread_id": thread_id, "assistant_id": assistant_id, "additional_instructions": additional_instructions, @@ -725,8 +726,8 @@ class AzureAssistantsAPI(BaseAzureLLM): "tools": tools, } if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return stream_fn(**base_data, event_handler=event_handler) + return stream_fn(**base_data) def run_thread_stream( self, diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 181894646e3..257ebf921d9 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -124,7 +124,7 @@ class OllamaChatConfig(BaseConfig): setattr(self.__class__, key, value) @classmethod - def get_config(cls): + def get_config(cls) -> dict[str, object]: return super().get_config() def get_supported_openai_params(self, model: str): diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index dccc83efed4..a9bbfedafd6 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -227,7 +227,7 @@ class OllamaConfig(BaseConfig): model: str, api_base: str | None = None, api_key: str | None = None, - ) -> Any: + ) -> dict[str, object] | None: """ curl http://localhost:11434/api/show -d '{ "name": "mistral" diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index edc8d64d9c2..15c543b94f0 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1,7 +1,7 @@ import time import types from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, cast import httpx @@ -2754,7 +2754,12 @@ class OpenAIAssistantsAPI(BaseLLM): message_thread: Final = await openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread( + id=message_thread.id, + created_at=message_thread.created_at, + metadata=message_thread.metadata, + object=message_thread.object, + ) # fmt: off @@ -2840,7 +2845,12 @@ class OpenAIAssistantsAPI(BaseLLM): message_thread: Final = openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread( + id=message_thread.id, + created_at=message_thread.created_at, + metadata=message_thread.metadata, + object=message_thread.object, + ) async def async_get_thread( self, @@ -2863,7 +2873,12 @@ class OpenAIAssistantsAPI(BaseLLM): response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread( + id=response.id, + created_at=response.created_at, + metadata=response.metadata, + object=response.object, + ) # fmt: off @@ -2929,7 +2944,12 @@ class OpenAIAssistantsAPI(BaseLLM): response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread( + id=response.id, + created_at=response.created_at, + metadata=response.metadata, + object=response.object, + ) def delete_thread(self): pass @@ -2986,18 +3006,27 @@ class OpenAIAssistantsAPI(BaseLLM): tools: Iterable[AssistantToolParam] | None, event_handler: AssistantEventHandler | None, ) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]: - data: Final[dict[str, Any]] = { - "thread_id": thread_id, - "assistant_id": assistant_id, - "additional_instructions": additional_instructions, - "instructions": instructions, - "metadata": metadata, - "model": model, - "tools": tools, - } + runs_stream: Final = client.beta.threads.runs.stream if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return runs_stream( + thread_id=thread_id, + assistant_id=assistant_id, + additional_instructions=additional_instructions, + instructions=instructions, + metadata=metadata, + model=model, + tools=tools, + event_handler=event_handler, + ) + return runs_stream( + thread_id=thread_id, + assistant_id=assistant_id, + additional_instructions=additional_instructions, + instructions=instructions, + metadata=metadata, + model=model, + tools=tools, + ) def run_thread_stream( self, diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py index 7ecc5e8ff3d..c79b6ffce43 100644 --- a/litellm/llms/vertex_ai/fine_tuning/handler.py +++ b/litellm/llms/vertex_ai/fine_tuning/handler.py @@ -280,7 +280,7 @@ class VertexFineTuningAPI(VertexLLM): vertex_location: str, vertex_credentials: str, request_route: str, - ): + ) -> object: _auth_header, vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, project_id=vertex_project, @@ -341,5 +341,4 @@ class VertexFineTuningAPI(VertexLLM): f"Error creating fine tuning job. Status code: {response.status_code}. Response: {response.text}" ) - response_json: Final = response.json() - return response_json + return response.json() diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index c6ad5928b74..fddc075bfc6 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -1,6 +1,7 @@ import base64 import json import os +from collections.abc import Mapping from io import BufferedRandom, BufferedReader, BytesIO from pathlib import Path from typing import TYPE_CHECKING, Any, Final, cast @@ -47,11 +48,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): image_edit_optional_params: ImageEditOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, Any]: + ) -> dict[str, object]: supported_params: Final = self.get_supported_openai_params(model) filtered_params = {key: value for key, value in image_edit_optional_params.items() if key in supported_params} - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} # Map OpenAI parameters to Imagen format if "n" in filtered_params: @@ -148,10 +149,10 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): model: str, prompt: str | None, image: FileTypes | None, - image_edit_optional_request_params: dict[str, Any], + image_edit_optional_request_params: Mapping[str, object], litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[dict[str, Any], RequestFiles | None]: + ) -> tuple[dict[str, object], RequestFiles | None]: # Prepare reference images in the correct Imagen format if image is None: raise ValueError("Vertex AI Imagen image edit requires at least one reference image.") @@ -182,14 +183,14 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): parameters["guidanceScale"] = 7.5 # Default guidance scale parameters["seed"] = None # Let Vertex AI choose random seed - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "instances": instances, "parameters": parameters, } - payload: Final[Any] = json.dumps(request_body) + payload: Final = json.dumps(request_body) empty_files: Final = cast(RequestFiles, []) - return cast(tuple[dict[str, Any], RequestFiles | None], (payload, empty_files)) + return cast(tuple[dict[str, object], RequestFiles | None], (payload, empty_files)) def transform_image_edit_response( self, @@ -237,8 +238,8 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): def _prepare_reference_images( self, image: FileTypes | list[FileTypes], - image_edit_optional_request_params: dict[str, Any], - ) -> list[dict[str, Any]]: + image_edit_optional_request_params: Mapping[str, object], + ) -> list[dict[str, object]]: """ Prepare reference images in the correct Imagen API format """ @@ -248,7 +249,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): else: images = [image] - reference_images: Final[list[dict[str, Any]]] = [] + reference_images: Final[list[dict[str, object]]] = [] for idx, img in enumerate(images): if img is None: @@ -258,7 +259,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): base64_data = base64.b64encode(image_bytes).decode("utf-8") # Create reference image structure - reference_image = { + reference_image: dict[str, object] = { "referenceType": "REFERENCE_TYPE_RAW", "referenceId": idx + 1, "referenceImage": {"bytesBase64Encoded": base64_data}, @@ -272,7 +273,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): mask_bytes: Final = self._read_all_bytes(mask_image) mask_base64: Final = base64.b64encode(mask_bytes).decode("utf-8") - mask_reference: Final = { + mask_reference: Final[dict[str, object]] = { "referenceType": "REFERENCE_TYPE_MASK", "referenceId": len(reference_images) + 1, "referenceImage": {"bytesBase64Encoded": mask_base64}, diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index d7a2491c04a..b2c52c53580 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -218,10 +218,10 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): contents: Final = [{"role": "user", "parts": [{"text": prompt}]}] # Prepare generation config - generation_config: Final[dict[str, Any]] = {"responseModalities": ["IMAGE"]} + generation_config: Final[dict[str, object]] = {"responseModalities": ["IMAGE"]} # Seed from user-supplied imageConfig dict; flat params are overlaid for backward compat. - image_config: Final[dict[str, Any]] = dict(optional_params.get("imageConfig") or {}) + image_config: Final[dict[str, object]] = dict(optional_params.get("imageConfig") or {}) if "aspectRatio" in optional_params: image_config["aspectRatio"] = optional_params["aspectRatio"] @@ -242,7 +242,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): elif "n" in optional_params: generation_config["candidateCount"] = optional_params["n"] - request_body: Final[dict[str, Any]] = { + request_body: Final[dict[str, object]] = { "contents": contents, "generationConfig": generation_config, } diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 082a90fdcfb..980a67b30d6 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -37,6 +37,7 @@ from litellm.repositories.table_repositories import ( MCPServerOAuthClientRepository, MCPServerRepository, MCPUserCredentialsRepository, + PrismaTableRepository, ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.verification_token_repository import ( @@ -522,11 +523,14 @@ def _user_credential_actions( return table +class _MCPUserEnvVarsRepository(PrismaTableRepository["prisma_db_models.LiteLLM_MCPUserEnvVars"]): + table_name = "litellm_mcpuserenvvars" + + def _user_env_var_actions( prisma_client: PrismaClient, ) -> "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]": - table: Final[TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = prisma_client.db.litellm_mcpuserenvvars - return table + return _MCPUserEnvVarsRepository(prisma_client).table async def _db_find_user_credential_row( diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index f2648c8466e..dc29a8377c9 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -48,7 +48,7 @@ from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManage from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository -from litellm.repositories.prisma_protocols import SpendLinkedTable +from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable from litellm.repositories.table_repositories import ( EndUserRepository, ModelAccessGroupBudgetRepository, @@ -435,6 +435,11 @@ class ResetBudgetJob: self.reset_settings: BudgetResetSettings = reset_settings or get_budget_reset_settings() self.pod_lock_manager: PodLockManager | None = pod_lock_manager + @property + def _new_batch(self) -> Callable[[], PrismaBatch]: + new_batch: Final[Callable[[], PrismaBatch]] = self.prisma_client.db.batch_ + return new_batch + async def _lease_is_held(self, lock_manager: PodLockManager) -> bool: """True only when the lease is readable and someone holds it. @@ -721,7 +726,7 @@ class ResetBudgetJob: ) async def _commit_budget_cascade_once(self, cascade: _BudgetCascade) -> None: - async with budget_cascade_unit_of_work(self.prisma_client.db.batch_) as uow: + async with budget_cascade_unit_of_work(self._new_batch) as uow: _queue_budget_linked_resets(uow.team_memberships, cascade) _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) @@ -861,7 +866,7 @@ class ResetBudgetJob: ) async def _write_key_reset_updates_once(self, updated_keys: list[LiteLLM_VerificationToken]) -> None: - async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: + async with spend_reset_unit_of_work(self._new_batch) as uow: for k in updated_keys: if k.token is None: continue @@ -885,7 +890,7 @@ class ResetBudgetJob: ) async def _write_user_reset_updates_once(self, updated_users: list[LiteLLM_UserTable]) -> None: - async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: + async with spend_reset_unit_of_work(self._new_batch) as uow: for u in updated_users: uow.users.queue_spend_reset( user_id=u.user_id, @@ -907,7 +912,7 @@ class ResetBudgetJob: ) async def _write_team_reset_updates_once(self, updated_teams: list[LiteLLM_TeamTable]) -> None: - async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: + async with spend_reset_unit_of_work(self._new_batch) as uow: for t in updated_teams: uow.teams.queue_spend_reset( team_id=t.team_id, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index facb822d00d..017ef6e09f6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal import httpx from fastapi import HTTPException +from typing_extensions import TypedDict, Unpack from litellm import DualCache from litellm._logging import verbose_proxy_logger @@ -111,6 +112,10 @@ class CiscoAIDefenseGuardrailAPIError(Exception): """Raised when there is an error talking to the Cisco AI Defense API.""" +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): """ Cisco AI Defense guardrail integration. @@ -144,7 +149,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): on_flagged_action: str | None = None, fallback_on_error: str | None = None, timeout: float | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: resolved_api_key: Final = api_key or os.environ.get("CISCO_AI_DEFENSE_API_KEY") if not resolved_api_key: diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index 214d4b486d4..539dc1ea1e9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py @@ -7,10 +7,10 @@ import os from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol import httpx -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version @@ -56,7 +56,13 @@ class DeepKeepFirewallResponse(TypedDict): class _DeepKeepInitKwargsView(TypedDict): """Typed read of the guardrail name carried in the untyped base-guardrail kwargs.""" - guardrail_name: ReadOnly[str] + guardrail_name: ReadOnly[str | None] + + +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + guardrail_name: ReadOnly[str | None] class _DeepKeepMetadataSource(TypedDict, total=False): @@ -110,7 +116,7 @@ class DeepKeepGuardrail(CustomGuardrail): firewall_id: str | None = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", extra_headers: Mapping[str, str] | list[str] | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index fa113aa4d33..eb62b896784 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -7,7 +7,7 @@ import time import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard +from typing import TYPE_CHECKING, ClassVar, Final, Literal, TypeGuard import httpx from fastapi import HTTPException @@ -50,6 +50,7 @@ from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.anthropic_messages.transformation import BaseAnthropicMessagesConfig from litellm.types.guardrails import LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -878,9 +879,9 @@ class HeadroomGuardrail(CustomGuardrail): async def async_pre_call_deployment_hook( self, - kwargs: dict[str, Any], + kwargs: dict[str, object], call_type: CallTypes | None, - ) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict + ) -> dict[str, object] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type) effective: Final = base_result if base_result is not None else kwargs if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES: @@ -897,7 +898,7 @@ class HeadroomGuardrail(CustomGuardrail): async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -919,8 +920,8 @@ class HeadroomGuardrail(CustomGuardrail): tools: dict, model: str, messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: BaseAnthropicMessagesConfig | None, anthropic_messages_optional_request_params: dict, logging_obj: LiteLLMLoggingObj | None, stream: bool, diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index fde40111d49..88c7f21ad72 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -360,7 +360,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): else: return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}} - def _should_block_content(self, armor_response: Mapping[str, Any], allow_sanitization: bool = False) -> bool: + def _should_block_content(self, armor_response: Mapping[str, object], allow_sanitization: bool = False) -> bool: """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult.""" for filt in self._filter_result_items(armor_response): # Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before @@ -429,7 +429,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return filter_results return [] - def _has_deidentify_match(self, armor_response: Mapping[str, Any]) -> bool: + def _has_deidentify_match(self, armor_response: Mapping[str, object]) -> bool: """Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction.""" for filter_entry in self._filter_result_items(armor_response): sdp = filter_entry.get("sdpFilterResult") @@ -439,7 +439,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): def _resolve_streaming_outcome( self, - armor_response: Mapping[str, Any], + armor_response: Mapping[str, object], assembled_response: object, content: str, ) -> tuple[bool, str | None]: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index bbc914a772a..12d23550cb9 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -8,7 +8,6 @@ POST /auto_router/validate_complexity_router_config - Dry-run the complexity-rou from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from itertools import chain, groupby -from operator import attrgetter from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Protocol from uuid import uuid4 @@ -1094,6 +1093,10 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]: ) +def _leg_group_id(leg: "_LegRow") -> str: + return leg.group_id + + class _LegRow(BaseModel): """One LiteLLM_ShadowEvalJob row, validated off the untyped prisma record. A row is one target's leg of a job; the legs of a job share group_id and identical config, @@ -1598,10 +1601,7 @@ async def list_shadow_eval_jobs( or () ) by_group: Final[Mapping[str, tuple[_LegRow, ...]]] = MappingProxyType( - { - group_id: tuple(group) - for group_id, group in groupby(sorted(legs, key=attrgetter("group_id")), key=attrgetter("group_id")) - } + {group_id: tuple(group) for group_id, group in groupby(sorted(legs, key=_leg_group_id), key=_leg_group_id)} ) newest_first: Final = sorted( by_group, key=lambda group_id: max(leg.created_at for leg in by_group[group_id]), reverse=True diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 0e6412a2c64..b8af432029f 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -31,11 +31,31 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: + from prisma import models as prisma_models + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient + +class _DailyTeamSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyTeamSpend"]): + table_name = "litellm_dailyteamspend" + + +def _daily_team_spend_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_DailyTeamSpend]": + """The sentinel rows this rollup writes, reads back and prunes.""" + return _DailyTeamSpendRepository(prisma_client).table + + +def _proxy_model_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_ProxyModelTable]": + """The stored deployments the rollup scans for PTU config.""" + return ModelRepository(prisma_client).table + + _HOURS_PER_DAY: Final = 24 _PRUNE_ID_CHUNK_SIZE: Final = 5_000 _UPSERT_ATTEMPTS: Final = 3 @@ -97,7 +117,7 @@ def _decode_model_info(raw: object) -> "Mapping[str, object] | None": """ if isinstance(raw, str): try: - decoded: Final = json.loads(raw) + decoded: Final[object] = json.loads(raw) except (TypeError, ValueError): return None return decoded if isinstance(decoded, dict) else None @@ -240,7 +260,7 @@ async def _upsert_ptu_daily_row( } } now: Final = datetime.now(timezone.utc) - await prisma_client.db.litellm_dailyteamspend.upsert( + await _daily_team_spend_table(prisma_client).upsert( where=where, data={ # mutable-ok: prisma upsert data payload "create": { # mutable-ok: prisma create payload @@ -353,7 +373,7 @@ async def _load_ptu_models(prisma_client: "PrismaClient", *, router: object | No The router is handed in rather than read off the proxy module, so a run prices exactly the deployments its caller declares and nothing a co-resident process left behind. """ - rows: Final = await prisma_client.db.litellm_proxymodeltable.find_many() + rows: Final = await _proxy_model_table(prisma_client).find_many() db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or ""))) config_records: Final = _config_deployments(router, owned_by_db=db_ids) models: Final = tuple( @@ -503,7 +523,7 @@ async def _existing_sentinel_keys( survives a rename. Nothing here reads the display name. """ date_range: Final = {"gte": start.isoformat(), "lte": end.isoformat()} # mutable-ok: prisma range filter - rows: Final = await prisma_client.db.litellm_dailyteamspend.find_many( + rows: Final = await _daily_team_spend_table(prisma_client).find_many( where={"api_key": PTU_SENTINEL_API_KEY, "date": date_range} # mutable-ok: prisma find filter ) return frozenset( @@ -771,7 +791,7 @@ async def _prune_unrefreshed_sentinel_rows( ) filters: Final = tuple(_prune_filter(date_str=date_str, cutoff=cutoff, chunk=chunk) for chunk in chunks) deletions: Final = tuple( - [await prisma_client.db.litellm_dailyteamspend.delete_many(where=where) for where in filters] + [await _daily_team_spend_table(prisma_client).delete_many(where=where) for where in filters] ) deleted: Final = sum(deletions) if deleted: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ddf31cb1d8a..21328a4962a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3401,7 +3401,7 @@ class _ConfigRow: __slots__ = ("param_name", "param_value") - def __init__(self, param_name: str, param_value: Any) -> None: + def __init__(self, param_name: str, param_value: object) -> None: self.param_name = param_name self.param_value = param_value @@ -3414,7 +3414,7 @@ def _pack_config_row(row: Any) -> dict[str, object]: return {"param_name": row.param_name, "param_value": row.param_value} -def _unpack_config_row(cached: Any) -> _ConfigRow | None: +def _unpack_config_row(cached: object) -> _ConfigRow | None: if cached is None or cached == _CONFIG_CACHE_MISS: return None if isinstance(cached, dict): @@ -3557,6 +3557,7 @@ class PrismaClient: verbose_proxy_logger.debug("Creating Prisma Client..") try: from prisma import Prisma + from prisma.types import DatasourceOverride except Exception as e: verbose_proxy_logger.error("Failed to import Prisma client: %s", e) verbose_proxy_logger.error("This usually means 'prisma generate' hasn't been run yet.") @@ -3607,11 +3608,11 @@ class PrismaClient: reader_token: Final = mint_database_token(token_auth, reader_iam_endpoint) read_replica_url = reader_iam_endpoint.build_url(reader_token) os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url - reader_kwargs: Final[dict[str, Any]] = {"datasource": {"url": read_replica_url}} + reader_datasource: Final = DatasourceOverride(url=read_replica_url) if http_client is not None: - reader_prisma = Prisma(http=http_client, **reader_kwargs) + reader_prisma = Prisma(http=http_client, datasource=reader_datasource) else: - reader_prisma = Prisma(**reader_kwargs) + reader_prisma = Prisma(datasource=reader_datasource) reader_wrapper: Final = PrismaWrapper( original_prisma=reader_prisma, token_auth=token_auth, diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index b824a5928c6..9d81c725c70 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -53,7 +53,7 @@ bedrock_realtime: Final = BedrockRealtime() xai_realtime: Final = XAIRealtime() vertex_llm_base: Final = VertexBase() base_llm_http_handler = BaseLLMHTTPHandler() -_EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_MODEL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]: diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index d24eb8ffc62..8ee76b93923 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -4,29 +4,23 @@ Model repository for database operations on LiteLLM_ProxyModelTable. import json from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final from litellm.models.model import LiteLLM_ProxyModelTable -from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) from litellm.repositories.base_repository import BaseRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models -class _PrismaModelDb(Protocol): - @property - def litellm_proxymodeltable(self) -> TableActions["prisma_models.LiteLLM_ProxyModelTable"]: ... - - -class _PrismaClientView(Protocol): - @property - def db(self) -> _PrismaModelDb: ... +class _ProxyModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_ProxyModelTable"]): + table_name = "litellm_proxymodeltable" class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): @@ -38,11 +32,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): @property def table(self) -> TableActions["prisma_models.LiteLLM_ProxyModelTable"]: - client: Final[_PrismaClientView] = self.prisma_client - return wrap_table_actions_for_config_sync( - actions=client.db.litellm_proxymodeltable, - table_name="litellm_proxymodeltable", - ) + return _ProxyModelTableRepository(self._prisma_client).table @property def model_class(self) -> type[LiteLLM_ProxyModelTable]: diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 5ff07d76b5d..cbe263699c9 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -5,6 +5,7 @@ Team repository for database operations on LiteLLM_TeamTable. import json from collections.abc import Mapping, Sequence from datetime import datetime +from types import TracebackType from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter @@ -40,6 +41,36 @@ def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays: return team +class _TeamTables(Protocol): + """The two team tables this repository reads and writes.""" + + @property + def litellm_teamtable(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]: ... + + @property + def litellm_deletedteamtable(self) -> TableActions["prisma_models.LiteLLM_DeletedTeamTable"]: ... + + +class _TeamTransactionManager(Protocol): + async def __aenter__(self) -> _TeamTables: ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> bool | None: ... + + +class _PrismaTeamDb(_TeamTables, Protocol): + def tx(self) -> _TeamTransactionManager: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _PrismaTeamDb: ... + + _MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member]) _JSON_ENCODED_TEAM_FIELDS: Final = ( "metadata", @@ -54,13 +85,18 @@ _JSON_ENCODED_TEAM_FIELDS: Final = ( class TeamRepository(BaseRepository[LiteLLM_TeamTable]): """Repository for team database operations.""" + @property + def _db(self) -> _PrismaTeamDb: + client: Final[_PrismaClientView] = self.prisma_client + return client.db + @property def table(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]: - return self.prisma_client.db.litellm_teamtable + return self._db.litellm_teamtable @property def deleted_table(self) -> TableActions["prisma_models.LiteLLM_DeletedTeamTable"]: - return self.prisma_client.db.litellm_deletedteamtable + return self._db.litellm_deletedteamtable @property def model_class(self) -> type[LiteLLM_TeamTable]: @@ -256,7 +292,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): archive_data["litellm_changed_by"] = litellm_changed_by archive_data["deleted_at"] = datetime.utcnow() - async with self.prisma_client.db.tx() as tx: + async with self._db.tx() as tx: await tx.litellm_deletedteamtable.create(data=archive_data) await tx.litellm_teamtable.delete(where={"team_id": team_id}) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 7fda5d96fb0..bf1838e50f7 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -266,7 +266,7 @@ def get_pre_routing_selection(kwargs: Mapping[str, object]) -> str | None: DISABLE_FALLBACKS_METADATA_KEY: Final = "_disable_fallbacks" -def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled: bool) -> None: +def record_disable_fallbacks(request_kwargs: Mapping[str, object] | None, disabled: bool) -> None: """ Write-or-clear the request's disable_fallbacks verdict into the router-internal metadata bucket. The wrapper pops the raw kwarg before any downstream frame runs, so the refusal @@ -286,7 +286,7 @@ def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled: bucket.pop(DISABLE_FALLBACKS_METADATA_KEY, None) -def fallbacks_disabled_for_request(kwargs: Mapping[str, Any]) -> bool: +def fallbacks_disabled_for_request(kwargs: Mapping[str, object]) -> bool: """True when this request opted out of fallbacks, read from the raw kwarg (pre-pop snapshots keep it) or the router-internal bucket the wrapper stamps after popping it.""" if kwargs.get("disable_fallbacks") is True: @@ -639,7 +639,7 @@ async def log_failure_fallback_event(original_model_group: str, kwargs: dict, or verbose_router_logger.error("Error in log_failure_fallback_event: %s", e) -def _check_non_standard_fallback_format(fallbacks: list[Any] | None) -> bool: +def _check_non_standard_fallback_format(fallbacks: Sequence[object] | None) -> bool: """ Checks if the fallbacks list is a list of strings or a list of dictionaries. @@ -653,8 +653,9 @@ def _check_non_standard_fallback_format(fallbacks: list[Any] | None) -> bool: return False if all(isinstance(item, str) for item in fallbacks): return True - elif all(isinstance(item, dict) for item in fallbacks): - for item in fallbacks: + dict_entries: Final = tuple(item for item in fallbacks if isinstance(item, dict)) + if len(dict_entries) == len(fallbacks): + for item in dict_entries: for key in LiteLLMParamsTypedDict.__annotations__: if key in item: # If the value is a list, it's likely a standard fallback model group mapping From f747bc67048f2b38e6e99bd57ca284fe6f77bdb5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:14:07 +0000 Subject: [PATCH 011/246] refactor(types): replace Any with real types across 29 more backend files Second batch of the fifth basedpyright Any reduction round. Every change is typing-only and leaves runtime behavior identical. Guardrail hooks and file, vector store and usage endpoints move their payload, header and response annotations from Any to object, Mapping[str, object] or the concrete response model the call site already produces. Two private aggregation helpers in the guardrail usage endpoints take a key accessor function instead of an attribute name string, so the key they read is checked against the row type. The verification token repository reaches its two tables through Protocols that name the handles it calls, rather than reading them off an untyped prisma client, and the Azure AD credential wrapper describes the azure-identity credential it wraps the same way. --- .../litellm_core_utils/llm_cost_calc/utils.py | 15 +++-- .../azure/text_to_speech/transformation.py | 9 +-- .../image_edit/stability_transformation.py | 4 +- .../responses/transformation.py | 14 ++--- .../llms/cohere/embed/v1_transformation.py | 15 +++-- litellm/llms/gdc/chat/transformation.py | 22 ++++++- .../llama3/transformation.py | 4 +- litellm/llms/voyage/rerank/transformation.py | 6 +- litellm/proxy/client/users.py | 3 +- .../proxy/common_utils/http_parsing_utils.py | 14 +++-- litellm/proxy/db/prisma_client.py | 9 ++- .../block_code_execution.py | 9 ++- .../cato_networks/cato_networks.py | 10 ++-- .../guardrail_hooks/compresr/compresr.py | 11 ++-- .../llm_as_a_judge/__init__.py | 23 +++++++- .../guardrails/guardrail_hooks/noma/noma.py | 15 +++-- .../panw_prisma_airs/panw_prisma_airs.py | 4 +- litellm/proxy/guardrails/usage_endpoints.py | 25 ++++---- .../proxy/hooks/proxy_track_cost_callback.py | 28 +++++++-- ...model_access_group_management_endpoints.py | 25 ++++---- .../usage_endpoints/ai_usage_chat.py | 11 ++-- .../openai_files_endpoints/files_endpoints.py | 4 +- .../storage_backend_service.py | 9 ++- .../proxy/vector_store_endpoints/endpoints.py | 2 +- .../vector_store_files_endpoints/endpoints.py | 10 ++-- litellm/proxy_auth/credentials.py | 58 ++++++++++++++----- litellm/rag/rag_query.py | 4 +- .../verification_token_repository.py | 46 ++++++++++++--- litellm/router_utils/search_api_router.py | 4 +- 29 files changed, 274 insertions(+), 139 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 68dc27ec25e..52dac92ee22 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone, tzinfo from types import MappingProxyType -from typing import Any, Final, Literal, TypedDict, cast +from typing import Final, Literal, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import litellm @@ -89,7 +89,7 @@ def _requested_image_size(optional_params: Mapping[str, object] | None) -> str | return value if value is not None and _IMAGE_SIZE_PATTERN.fullmatch(value) else None -def get_web_search_requests(server_tool_use: Any) -> int | None: +def get_web_search_requests(server_tool_use: object) -> int | None: """ Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance, @@ -1494,7 +1494,7 @@ def calculate_image_response_cost_from_usage( if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0: return None - input_tokens_details: Final = getattr(usage, "input_tokens_details", None) + input_tokens_details: Final[object] = getattr(usage, "input_tokens_details", None) prompt_tokens_details: PromptTokensDetailsWrapper | None = None if input_tokens_details is not None: # input_tokens_details may be a dict (e.g. OpenAI image edit responses) @@ -1507,9 +1507,12 @@ def calculate_image_response_cost_from_usage( cached_tokens=0, ) - output_tokens_details = getattr(usage, "completion_tokens_details", None) - if output_tokens_details is None: - output_tokens_details = getattr(usage, "output_tokens_details", None) + completion_tokens_details_attr: Final[object] = getattr(usage, "completion_tokens_details", None) + output_tokens_details: Final[object] = ( + getattr(usage, "output_tokens_details", None) + if completion_tokens_details_attr is None + else completion_tokens_details_attr + ) if output_tokens_details is None: completion_tokens_details = CompletionTokensDetailsWrapper( diff --git a/litellm/llms/azure/text_to_speech/transformation.py b/litellm/llms/azure/text_to_speech/transformation.py index d8ccf26ce60..eed7a3178ca 100644 --- a/litellm/llms/azure/text_to_speech/transformation.py +++ b/litellm/llms/azure/text_to_speech/transformation.py @@ -19,6 +19,7 @@ from litellm.secret_managers.main import get_secret_str if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.openai import HttpxBinaryResponseContent else: LiteLLMLoggingObj = Any @@ -67,15 +68,15 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): litellm_params_dict: dict, logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout, - extra_headers: dict[str, Any] | None, - base_llm_http_handler: Any, + extra_headers: dict[str, object] | None, + base_llm_http_handler: "BaseLLMHTTPHandler", aspeech: bool, api_base: str | None, api_key: str | None, - **kwargs: Any, + **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[object, object, "HttpxBinaryResponseContent"], ]: """ Dispatch method to handle Azure AVA TTS requests diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index bc9a64f587a..01e25f4671e 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -125,7 +125,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): } # Create a copy to not mutate original - convert TypedDict to regular dict - mapped_params: Final[dict[str, Any]] = dict(image_edit_optional_params) + mapped_params: Final[dict[str, object]] = dict(image_edit_optional_params) for k, v in image_edit_optional_params.items(): if k in param_mapping: @@ -172,7 +172,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): Returns the request body dict that will be JSON-encoded by the handler. """ # Build Bedrock Stability request - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "output_format": "png", # Default to PNG } diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index bbbda4d14b6..3d2eab8fcee 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -16,8 +16,8 @@ BaseAWSLLM._sign_request after the request body is finalized. """ import json -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final import httpx from typing_extensions import ReadOnly, TypedDict @@ -142,9 +142,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI return False @staticmethod - def _filter_unsupported_tools(tools: list[Any]) -> list[Any]: + def _filter_unsupported_tools(tools: "Sequence[object]") -> "list[object]": """Keep only tool types Mantle's Responses API accepts.""" - kept: Final[list[Any]] = [] + kept: Final[list[object]] = [] dropped_types: Final[list[str]] = [] for tool in tools: if not isinstance(tool, dict): @@ -217,11 +217,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) @staticmethod - def _is_codex_additional_tools_item(item: Any) -> bool: + def _is_codex_additional_tools_item(item: object) -> bool: return isinstance(item, dict) and item.get("type") == _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE @staticmethod - def _tools_of_additional_tools_item(item: "dict[str, Any]") -> "list[Any]": + def _tools_of_additional_tools_item(item: "Mapping[str, object]") -> "list[object]": tools: Final = item.get("tools") return tools if isinstance(tools, list) else [] @@ -229,7 +229,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def _hoist_codex_additional_tools( cls, input: "str | ResponseInputParam", - ) -> "tuple[str | ResponseInputParam, list[Any]]": + ) -> "tuple[str | ResponseInputParam, list[object]]": """Codex's "responses lite" wire mode ships tool definitions inside `input` as {"type": "additional_tools", "role": "developer", "tools": [...]} items. api.openai.com accepts that item type; Mantle diff --git a/litellm/llms/cohere/embed/v1_transformation.py b/litellm/llms/cohere/embed/v1_transformation.py index ee40464362d..b35fae5a1ac 100644 --- a/litellm/llms/cohere/embed/v1_transformation.py +++ b/litellm/llms/cohere/embed/v1_transformation.py @@ -2,7 +2,8 @@ Legacy /v1/embedding transformation logic for Bedrock Cohere. """ -from typing import Any, Final +from collections.abc import Sized +from typing import Final, Protocol import httpx @@ -16,6 +17,12 @@ from litellm.types.utils import EmbeddingResponse, PromptTokensDetailsWrapper, U from litellm.utils import is_base64_encoded +class _SupportsEncode(Protocol): + """Tokenizer handle: the embedding usage path only encodes text to measure its token length.""" + + def encode(self, text: str, /) -> Sized: ... + + class CohereEmbeddingConfig: """ Reference: https://docs.cohere.com/v2/reference/embed @@ -61,7 +68,7 @@ class CohereEmbeddingConfig: return transformed_request - def _calculate_usage(self, input: list[str], encoding: Any, meta: dict) -> Usage: + def _calculate_usage(self, input: list[str], encoding: _SupportsEncode, meta: dict) -> Usage: input_tokens = 0 text_tokens: Final[int | None] = meta.get("billed_units", {}).get("input_tokens") @@ -97,7 +104,7 @@ class CohereEmbeddingConfig: data: dict | CohereEmbeddingRequest, model_response: EmbeddingResponse, model: str, - encoding: Any, + encoding: _SupportsEncode, input: list, ) -> EmbeddingResponse: response_json: Final = response.json() @@ -121,7 +128,7 @@ class CohereEmbeddingConfig: response_json: dict, model_response: EmbeddingResponse, model: str, - encoding: Any, + encoding: _SupportsEncode, input: list, ) -> EmbeddingResponse: """ diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index 03037512551..6eac3ac79cd 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -7,14 +7,32 @@ import os import re import threading from collections.abc import Callable -from typing import Any, Final, Protocol +from typing import Final, Protocol from urllib.parse import urlsplit +from typing_extensions import ReadOnly, TypedDict, Unpack + import litellm from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.types.llms.openai import AllMessageValues +class _OpenAIGPTConfigOptions(TypedDict, total=False): + """The sampling defaults ``OpenAIGPTConfig.__init__`` accepts and stashes on the class.""" + + frequency_penalty: ReadOnly[int | None] + function_call: ReadOnly[str | dict[str, object] | None] + functions: ReadOnly[list[object] | None] + logit_bias: ReadOnly[dict[str, object] | None] + max_tokens: ReadOnly[int | None] + n: ReadOnly[int | None] + presence_penalty: ReadOnly[int | None] + stop: ReadOnly[str | list[object] | None] + temperature: ReadOnly[int | None] + top_p: ReadOnly[int | None] + response_format: ReadOnly[dict[str, object] | None] + + class _GDCHAudienceCredentials(Protocol): """A GDCH service account credential already bound to an audience, ready to mint a bearer token.""" @@ -32,7 +50,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): _GDCH_CREDENTIAL_TYPE: Final[str] = "gdch_service_account" _PATH_ID_PATTERN: Final[re.Pattern[str]] = re.compile(r"^[a-zA-Z0-9_-]+$") - def __init__(self, **kwargs: Any) -> None: + def __init__(self, **kwargs: Unpack[_OpenAIGPTConfigOptions]) -> None: super().__init__(**kwargs) self._creds_lock = threading.Lock() self._gdch_creds_cache: dict[tuple[str, str], _GDCHAudienceCredentials] = {} diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py index 279035c455d..89a5b8a570e 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py @@ -1,6 +1,6 @@ import types from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -95,7 +95,7 @@ class VertexAILlama3Config(OpenAIGPTConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "VertexAILlama3StreamingHandler": return VertexAILlama3StreamingHandler( streaming_response=streaming_response, sync_stream=sync_stream, diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index fea8452d934..6f8d024f0b1 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -4,8 +4,8 @@ Transformation logic for Voyage AI's /v1/rerank endpoint. Docs - https://docs.voyageai.com/docs/reranker """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final import httpx @@ -33,7 +33,7 @@ class VoyageRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: Sequence[str | Mapping[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index 3f11fe94043..503c92228a8 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final import requests @@ -50,7 +51,7 @@ class UsersManagementClient: response.raise_for_status() return response.json() - def create_user(self, user_data: dict[str, Any]) -> dict[str, Any]: + def create_user(self, user_data: Mapping[str, object]) -> dict[str, Any]: """Create a new user (POST /user/new)""" url: Final = f"{self.base_url}/user/new" response: Final = requests.post(url, headers=self._get_headers(), json=user_data, timeout=self.timeout) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 54a0f18fd63..2bae3e946f7 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -52,14 +52,18 @@ def _unqualified(annotation: object) -> object: return _unqualified(qualified[0]) +def _union_members(annotation: object) -> tuple[object, ...]: + """The non-``None`` members of a union annotation, or the annotation itself when it is not a union.""" + if get_origin(annotation) not in (Union, UnionType): + return (annotation,) + members: Final[tuple[object, ...]] = get_args(annotation) + return tuple(arg for arg in members if arg is not type(None)) + + def _numeric_form_type(annotation: object) -> type[int] | type[float] | None: """The scalar to parse an ``int``/``float``-typed field as, else ``None``.""" unwrapped: Final = _unqualified(annotation) - candidates: Final = ( - tuple(arg for arg in get_args(unwrapped) if arg is not type(None)) - if get_origin(unwrapped) in (Union, UnionType) - else (unwrapped,) - ) + candidates: Final = _union_members(unwrapped) if len(candidates) != 1: return None if candidates[0] is int: diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 2190ae55fd2..2f2ebfdf2bb 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -13,7 +13,7 @@ import urllib import urllib.parse from collections.abc import Callable from datetime import datetime, timedelta -from typing import Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_proxy_logger from litellm.proxy.db.token_auth import ( @@ -27,6 +27,9 @@ from litellm.proxy.db.token_auth import ( ) from litellm.secret_managers.main import str_to_bool +if TYPE_CHECKING: + from prisma import Prisma + __all__ = ( "IAMEndpoint", "PrismaManager", @@ -242,7 +245,7 @@ class PrismaWrapper: def _write_engine(prisma_client: _PrismaClient, engine: _PrismaEngine) -> None: prisma_client._Prisma__engine = engine - def _instrument_prisma_client(self, prisma_client: _PrismaClient) -> _PrismaDrainTracker | None: + def _instrument_prisma_client(self, prisma_client: "Prisma | _PrismaClient") -> _PrismaDrainTracker | None: from prisma.errors import ClientNotConnectedError try: @@ -255,7 +258,7 @@ class PrismaWrapper: self._write_engine(prisma_client, _TrackedPrismaEngine(engine, tracker)) return tracker - def _get_engine_pid(self, prisma_client: _PrismaClient | None = None) -> int: + def _get_engine_pid(self, prisma_client: "Prisma | _PrismaClient | None" = None) -> int: """Get the PID of the current Prisma engine subprocess, or 0 if unavailable. Must never raise: it runs inside the reconnect path, where the client diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py index bf2aa1f76e0..2b697671eda 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -8,9 +8,10 @@ confidence scoring and a tunable threshold (only block when confidence >= thresh import re from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, cast from fastapi import HTTPException +from typing_extensions import TypedDict, Unpack from litellm.integrations.custom_guardrail import ( CustomGuardrail, @@ -314,6 +315,10 @@ def _confidence_for_block( return 0.0 +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + class BlockCodeExecutionGuardrail(CustomGuardrail): """ Guardrail that detects fenced code blocks (markdown ```) and blocks or masks them @@ -332,7 +337,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): detect_execution_intent: bool = True, event_hook: Literal["pre_call", "post_call", "during_call"] | list[str] | None = None, default_on: bool = False, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: # Normalize to type expected by CustomGuardrail _event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 176c308eda6..2d203c31974 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -264,7 +264,7 @@ class CatoNetworksGuardrail(CustomGuardrail): stack.extend(reversed(node)) @classmethod - def _extra_inspection_sources(cls, data: Mapping[str, Any]) -> Sequence[tuple[str, Sequence[Mapping[str, str]]]]: + def _extra_inspection_sources(cls, data: Mapping[str, object]) -> Sequence[tuple[str, Sequence[Mapping[str, str]]]]: """Text the proxy forwards to the model outside chat ``messages``: Responses-API ``input`` and ``instructions``, legacy completion ``prompt`` and tool/function/``response_format`` schema strings. Returned @@ -336,7 +336,7 @@ class CatoNetworksGuardrail(CustomGuardrail): ) raise HTTPException(status_code=400, detail=detection_message) - def _anonymize_request(self, res: Any, data: dict) -> dict: + def _anonymize_request(self, res: _CatoAnalyzeResponse, data: dict) -> dict: verbose_proxy_logger.info("Cato: anonymize action") redacted_chat: Final = res.get("redacted_chat") if not redacted_chat: @@ -379,7 +379,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return data @classmethod - def _apply_extra_redaction(cls, data: dict, field: str, redacted: list) -> bool: + def _apply_extra_redaction(cls, data: dict, field: str, redacted: Sequence[Mapping[str, object]]) -> bool: if field == "input": input_only: Final = {"input": data["input"]} if not redacted: @@ -400,7 +400,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return True @classmethod - def _apply_schema_string_redaction(cls, data: dict, redacted: list) -> None: + def _apply_schema_string_redaction(cls, data: dict, redacted: Sequence[Mapping[str, object]]) -> None: redactions: Final = iter(redacted) for container, key in cls._iter_schema_string_refs(data): replacement = next(redactions, None) @@ -408,7 +408,7 @@ class CatoNetworksGuardrail(CustomGuardrail): container[key] = replacement["content"] @staticmethod - def _apply_prompt_redaction(data: dict, redacted: list) -> None: + def _apply_prompt_redaction(data: dict, redacted: Sequence[Mapping[str, object]]) -> None: contents: Final = [m.get("content") for m in redacted if isinstance(m, dict)] prompt: Final = data.get("prompt") if isinstance(prompt, str): diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 93d859066b0..1ecdb1b0f63 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -22,7 +22,7 @@ import json import time from collections import Counter, OrderedDict from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Final, Literal, TypeGuard +from typing import TYPE_CHECKING, Final, Literal, TypeGuard from urllib.parse import urlparse import httpx @@ -64,6 +64,9 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) from litellm.types.proxy.guardrails.guardrail_hooks.base import ( GuardrailConfigModel, ) @@ -1049,7 +1052,7 @@ class CompresrGuardrail(CustomGuardrail): async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -1069,8 +1072,8 @@ class CompresrGuardrail(CustomGuardrail): tools: dict, model: str, messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: BaseAnthropicMessagesConfig | None, anthropic_messages_optional_request_params: dict, logging_obj: LiteLLMLoggingObj | None, stream: bool, diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index 172b1440ca3..806ab5161e8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -2,10 +2,10 @@ from collections.abc import Callable, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeVar +from typing import TYPE_CHECKING, Final, Generic, Literal, Optional, TypeVar from fastapi import HTTPException -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -106,6 +106,23 @@ def _build_judge_prompt( ) +class _CustomGuardrailOptions(TypedDict, total=False): + """The ``CustomGuardrail`` options this guardrail accepts and forwards untouched.""" + + mask_request_content: ReadOnly[bool] + mask_response_content: ReadOnly[bool] + violation_message_template: ReadOnly[str | None] + end_session_after_n_fails: ReadOnly[int | None] + on_violation: ReadOnly[str | None] + realtime_violation_message: ReadOnly[str | None] + on_sensitive_data: ReadOnly[str | None] + sensitive_data_route_to_model: ReadOnly[str | None] + sticky_session_routing: ReadOnly[bool] + run_in_parallel: ReadOnly[bool] + scan_raw_request: ReadOnly[bool] + only_scan_new_messages: ReadOnly[bool] + + class LLMAsAJudgeGuardrail(CustomGuardrail): """Post-call guardrail that judges response quality via an LLM.""" @@ -119,7 +136,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, default_on: bool = False, router_provider: "Callable[[], Router | None] | None" = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: _event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None if event_hook is not None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 7ef0a9f73f3..edd78e0bbc6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -9,13 +9,13 @@ import asyncio import json import os import warnings -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterable from datetime import datetime from typing import ( TYPE_CHECKING, - Any, Final, Literal, + TypeVar, ) from urllib.parse import urljoin @@ -39,9 +39,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( CallTypes, CallTypesLiteral, - EmbeddingResponse, GuardrailStatus, - ImageResponse, ModelResponseStream, TextCompletionResponse, ) @@ -53,7 +51,8 @@ SENSITIVE_DATA_DETECTOR_KEYS: Final[list[str]] = ["sensitiveData", "dataDetector # Type aliases MessageRole = Literal["user", "assistant"] -LLMResponse = Any | ModelResponse | EmbeddingResponse | ImageResponse +LLMResponse = object +_LLMResponseT: Final = TypeVar("_LLMResponseT") _LEGACY_NOMA_DEPRECATION_WARNED = False if TYPE_CHECKING: @@ -709,10 +708,10 @@ class NomaGuardrail(CustomGuardrail): async def _check_llm_response( self, request_data: dict, - response: LLMResponse, + response: _LLMResponseT, user_auth: UserAPIKeyAuth, event_type: GuardrailEventHooks | None = None, - ) -> Any: + ) -> _LLMResponseT: """Check LLM response for policy violations""" content: Final = await self._process_llm_response_check(request_data, response, user_auth, event_type) if not content: @@ -798,7 +797,7 @@ class NomaGuardrail(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[ModelResponseStream], request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: """Process streaming response chunks with Noma guardrail.""" diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index b73d3adb99e..86ad2f9db5f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -793,7 +793,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): }, ) - def _prepare_metadata_from_request(self, data: dict[str, Any]) -> dict[str, Any]: + def _prepare_metadata_from_request(self, data: dict[str, Any]) -> dict[str, object]: """ Extract and prepare metadata from request data for PANW API call. @@ -809,7 +809,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): """ user_metadata: Final = data.get("metadata", {}) or {} requester_meta: Final = user_metadata.get("requester_metadata", {}) or {} - metadata: Final = { + metadata: Final[dict[str, object]] = { "user": data.get("user") or "litellm_user", "model": data.get("model") or "unknown", } diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 6259efb6654..3f7c36bbbf2 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -45,6 +45,7 @@ _EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({}) _ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"passed": 0, "flagged": 1, "blocked": 2}) _T = TypeVar("_T") +_MetricsRowT = TypeVar("_MetricsRowT", bound="_DailyMetricsRow") _USAGE_MAX_RANGE_DAYS: Final = 366 @@ -360,10 +361,12 @@ def _trend_from_comparison(current_fail: float, previous_fail: float) -> str: return "stable" -def _aggregate_daily_metrics(metrics: "Sequence[_DailyMetricsRow]", id_attr: str) -> Mapping[str, _MetricTotals]: +def _aggregate_daily_metrics( + metrics: "Sequence[_MetricsRowT]", id_of: "Callable[[_MetricsRowT], str]" +) -> Mapping[str, _MetricTotals]: agg: Final[dict[str, _MetricTotals]] = {} for m in metrics: - gid: str = getattr(m, id_attr) + gid: str = id_of(m) if gid not in agg: agg[gid] = {"requests": 0, "passed": 0, "blocked": 0, "flagged": 0} agg[gid]["requests"] += int(m.requests_evaluated or 0) @@ -373,10 +376,12 @@ def _aggregate_daily_metrics(metrics: "Sequence[_DailyMetricsRow]", id_attr: str return agg -def _prev_fail_rates(metrics_prev: "Sequence[_DailyMetricsRow]", id_attr: str) -> Mapping[str, float]: +def _prev_fail_rates( + metrics_prev: "Sequence[_MetricsRowT]", id_of: "Callable[[_MetricsRowT], str]" +) -> Mapping[str, float]: prev_agg_raw: Final[dict[str, _PrevPeriodCounts]] = {} for m in metrics_prev: - gid: str = getattr(m, id_attr) + gid: str = id_of(m) r, b = int(m.requests_evaluated or 0), int(m.blocked_count or 0) if gid not in prev_agg_raw: prev_agg_raw[gid] = {"req": 0, "blocked": 0} @@ -429,7 +434,7 @@ def _field_str(mapping: Mapping[str, object], key: str, default: str) -> str: return str(mapping.get(key, default)) -def _get_guardrail_attrs(g: "_DbOrConfigGuardrail") -> tuple[Any, str]: +def _get_guardrail_attrs(g: "_DbOrConfigGuardrail") -> tuple[str | None, str]: """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" gid: Final = _get_guardrail_field(g, "guardrail_id") name: Final = _get_guardrail_field(g, "guardrail_name") @@ -592,8 +597,8 @@ async def guardrails_usage_overview( Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits] ] = await _find_daily_guardrail_usage_units(prisma_client, where=units_where) - agg: Final = _aggregate_daily_metrics(metrics, "guardrail_id") - prev_agg: Final = _prev_fail_rates(metrics_prev, "guardrail_id") + agg: Final = _aggregate_daily_metrics(metrics, lambda m: m.guardrail_id) + prev_agg: Final = _prev_fail_rates(metrics_prev, lambda m: m.guardrail_id) units_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_counter_units) cost_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_tracked_cost) untracked_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_untracked_units) @@ -811,7 +816,7 @@ def _usage_log_entry_from_row( ) -def _snippet(text: Any, max_len: int = 200) -> str | None: +def _snippet(text: object, max_len: int = 200) -> str | None: if text is None: return None if isinstance(text, str): @@ -964,8 +969,8 @@ async def policies_usage_overview( } }, ) - agg: Final = _aggregate_daily_metrics(metrics, "policy_id") - prev_agg: Final = _prev_fail_rates(metrics_prev, "policy_id") + agg: Final = _aggregate_daily_metrics(metrics, lambda m: m.policy_id) + prev_agg: Final = _prev_fail_rates(metrics_prev, lambda m: m.policy_id) chart: Final = _chart_from_metrics(metrics) total_requests: Final = sum(a["requests"] for a in agg.values()) total_blocked: Final = sum(a["blocked"] for a in agg.values()) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c4fba8ecf9e..e6ada41f062 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -2,7 +2,7 @@ import asyncio import traceback from collections.abc import Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm from litellm._logging import verbose_proxy_logger @@ -578,18 +578,36 @@ def _get_request_tags_for_cost_tracking( return None +class _IncrementSpendCounters(Protocol): + """The ``increment_spend_counters`` coroutine :func:`_update_database_and_spend_counters` awaits.""" + + async def __call__( + self, + token: str | None, + team_id: str | None, + user_id: str | None, + response_cost: float | None, + org_id: str | None = None, + budget_reservation: dict[str, object] | None = None, + end_user_id: str | None = None, + tags: list[str] | None = None, + request_started_at: datetime | None = None, + model_access_groups: Sequence[str] | None = None, + ) -> None: ... + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", - increment_spend_counters: Any, + increment_spend_counters: _IncrementSpendCounters, user_api_key: str | None, user_id: str | None, end_user_id: str | None, team_id: str | None, org_id: str | None, kwargs: dict, - completion_response: litellm.ModelResponse | Any | None, - start_time: Any, - end_time: Any, + completion_response: object, + start_time: datetime | None, + end_time: datetime | None, response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index a48130a4f22..e960bdfe337 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -540,7 +540,7 @@ async def get_all_access_groups_from_db( deployments: Final = await ModelRepository(prisma_client).table.find_many() # Build access group map - access_group_map: Final[dict[str, dict[str, Any]]] = {} + model_names_by_group: Final[dict[str, list[str]]] = {} for deployment in deployments: model_info = deployment.model_info or {} @@ -550,25 +550,20 @@ async def get_all_access_groups_from_db( model_name = deployment.model_name for access_group in access_groups: - if access_group not in access_group_map: - access_group_map[access_group] = { - "model_names": set(), - "deployment_count": 0, - } + if access_group not in model_names_by_group: + model_names_by_group[access_group] = [] - access_group_map[access_group]["model_names"].add(model_name) - access_group_map[access_group]["deployment_count"] += 1 + model_names_by_group[access_group].append(model_name) # Convert to AccessGroupInfo objects - result: Final = {} - for access_group, data in access_group_map.items(): - result[access_group] = AccessGroupInfo( + return { + access_group: AccessGroupInfo( access_group=access_group, - model_names=sorted(list(data["model_names"])), - deployment_count=data["deployment_count"], + model_names=sorted(frozenset(model_names)), + deployment_count=len(model_names), ) - - return result + for access_group, model_names in model_names_by_group.items() + } @router.post( diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index da4ddbd0aac..1265da99d89 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -6,7 +6,7 @@ usage/spend data by querying the aggregated daily activity endpoints. import json from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import date -from typing import Any, Final, Literal, NamedTuple, Protocol, cast, overload +from typing import Final, Literal, NamedTuple, Protocol, cast, overload from typing_extensions import ReadOnly, TypedDict @@ -16,6 +16,7 @@ from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) +from litellm.types.utils import ChatCompletionMessageToolCall # --------------------------------------------------------------------------- # Constants @@ -489,19 +490,19 @@ async def _execute_tool_call( async def _process_tool_call( - tc: Any, + tc: ChatCompletionMessageToolCall, chat_messages: list[Mapping[str, object]], user_id: str | None, is_admin: bool, ) -> AsyncIterator[str]: """Execute a single tool call, yielding SSE events for status.""" - fn_name: Final[str] = tc.function.name + fn_name: Final = tc.function.name fn_args: Final[Mapping[str, str]] = json.loads(tc.function.arguments) allowed_names: Final = {t["function"]["name"] for t in get_tools_for_role(is_admin)} - handler: Final = TOOL_HANDLERS.get(fn_name) + handler: Final = TOOL_HANDLERS.get(fn_name) if fn_name is not None else None - if fn_name not in allowed_names or not handler: + if fn_name is None or fn_name not in allowed_names or not handler: chat_messages.append( { "role": "tool", diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index bf07f4748ef..4d3a397e519 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -485,8 +485,8 @@ async def create_file( # Parse expires_after if provided expires_after: FileExpiresAfter | None = None form_data_raw: Final = await request.form() - form_data_dict: Final[dict[str, Any]] = dict(form_data_raw) - extracted_litellm_metadata: Final[dict[str, Any] | None] = extract_nested_form_metadata( + form_data_dict: Final[Mapping[str, object]] = dict(form_data_raw) + extracted_litellm_metadata: Final[Mapping[str, object] | None] = extract_nested_form_metadata( form_data=form_data_dict, prefix="litellm_metadata[" ) expires_after_anchor: Final = form_data_raw.get("expires_after[anchor]") diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index e766f335071..0407499cbcc 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -7,7 +7,6 @@ storage backends (e.g., Azure Blob Storage) and managing associated metadata. import base64 import time -from collections.abc import Mapping from typing import Any, Final, cast from litellm._logging import verbose_proxy_logger @@ -17,7 +16,7 @@ from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.llms.openai import OpenAIFileObject, OpenAIFilesPurpose -from litellm.types.utils import SpecialEnums +from litellm.types.utils import ExtractedFileData, SpecialEnums class StorageBackendFileService: @@ -33,7 +32,7 @@ class StorageBackendFileService: @staticmethod async def upload_file_to_storage_backend( - file_data: Mapping[str, Any], + file_data: ExtractedFileData, target_storage: str, target_model_names: list[str], purpose: OpenAIFilesPurpose, @@ -163,7 +162,7 @@ class StorageBackendFileService: @staticmethod def _create_unified_file_id( - file_type: str, + file_type: str | None, target_model_names: list[str], file_id: str, ) -> str: @@ -193,7 +192,7 @@ class StorageBackendFileService: @staticmethod async def _store_in_managed_files( file_object: OpenAIFileObject, - file_data: Mapping[str, Any], + file_data: ExtractedFileData, target_model_names: list[str], target_storage: str, storage_url: str, diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 1feda0b0bb5..f21c294e5a2 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -244,7 +244,7 @@ async def vector_store_create( ) # Create vector store across multiple models - response: Final = await managed_vector_stores.acreate_vector_store( + response: Final[object] = await managed_vector_stores.acreate_vector_store( create_request=data, llm_router=llm_router, target_model_names_list=target_model_names_list, diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 957ed9fd0b9..3ddea288ab8 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -553,7 +553,7 @@ async def vector_store_file_create( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -756,7 +756,7 @@ async def vector_store_file_retrieve( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -863,7 +863,7 @@ async def vector_store_file_content( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -973,7 +973,7 @@ async def vector_store_file_update( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -1080,7 +1080,7 @@ async def vector_store_file_delete( processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response = await processor.base_process_llm_request( + response: object = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy_auth/credentials.py b/litellm/proxy_auth/credentials.py index a4e29241959..f8814954a7a 100644 --- a/litellm/proxy_auth/credentials.py +++ b/litellm/proxy_auth/credentials.py @@ -7,7 +7,7 @@ It follows the same TokenCredential protocol used by Azure SDK. import time from dataclasses import dataclass -from typing import Any, Final, Protocol, runtime_checkable +from typing import Final, Protocol, runtime_checkable @dataclass @@ -50,6 +50,22 @@ class TokenCredential(Protocol): ... +class _AzureAccessToken(Protocol): + """The two attributes :class:`AzureADCredential` reads off an azure-identity token.""" + + @property + def token(self) -> str: ... + + @property + def expires_on(self) -> int: ... + + +class _AzureTokenCredential(Protocol): + """The single method :class:`AzureADCredential` calls on the credential it wraps.""" + + def get_token(self, *scopes: str) -> _AzureAccessToken: ... + + class AzureADCredential: """ Wrapper for Azure Identity credentials. @@ -71,7 +87,7 @@ class AzureADCredential: cred = AzureADCredential(credential=azure_cred) """ - def __init__(self, credential: Any | None = None): + def __init__(self, credential: _AzureTokenCredential | None = None): """ Initialize with an optional Azure credential. @@ -79,7 +95,7 @@ class AzureADCredential: credential: An azure-identity credential object. If None, DefaultAzureCredential will be used on first token request. """ - self._credential: Any = credential + self._credential: _AzureTokenCredential | None = credential self._initialized = credential is not None def get_token(self, scope: str) -> AccessToken: @@ -95,20 +111,30 @@ class AzureADCredential: Raises: ImportError: If azure-identity is not installed. """ - if not self._initialized: - try: - from azure.identity import DefaultAzureCredential - - self._credential = DefaultAzureCredential() - self._initialized = True - except ImportError: - raise ImportError( - "azure-identity is required for AzureADCredential. Install it with: pip install azure-identity" - ) - - result: Final = self._credential.get_token(scope) + result: Final = self._resolve_credential().get_token(scope) return AccessToken(token=result.token, expires_on=result.expires_on) + def _resolve_credential(self) -> _AzureTokenCredential: + """Return the wrapped credential, building the Azure default chain on first use. + + Raises: + ImportError: If azure-identity is not installed. + """ + existing: Final = self._credential + if existing is not None: + return existing + try: + from azure.identity import DefaultAzureCredential + + created: Final = DefaultAzureCredential() + except ImportError: + raise ImportError( + "azure-identity is required for AzureADCredential. Install it with: pip install azure-identity" + ) + self._credential = created + self._initialized = True + return created + class GenericOAuth2Credential: """ @@ -228,7 +254,7 @@ class ProxyAuthHandler: self._cached_token = self.credential.get_token(self.scope) return self._cached_token - def get_auth_headers(self) -> dict: + def get_auth_headers(self) -> dict[str, str]: """ Get HTTP headers for authentication. diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index 255faf94402..9325547c17d 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -124,9 +124,9 @@ class RAGQuery: @staticmethod def extract_documents_from_search( search_response: Any, - ) -> list[str | dict[str, Any]]: + ) -> list[str | dict[str, object]]: """Extract text documents from vector store search response.""" - documents: Final[list[str | dict[str, Any]]] = [] + documents: Final[list[str | dict[str, object]]] = [] search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])} for result in search_data["results"]: content_list = result.get("content", []) diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index c0e59f9b975..d02c2114136 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -5,7 +5,8 @@ VerificationToken repository for database operations on LiteLLM_VerificationToke import json from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Final +from types import TracebackType +from typing import TYPE_CHECKING, Final, Protocol from litellm.models.verification_token import ( LiteLLM_VerificationToken, @@ -25,7 +26,36 @@ if TYPE_CHECKING: LiteLLM_VerificationToken as PrismaVerificationToken, ) - from litellm.proxy.utils import PrismaClient + +class _VerificationTokenTables(Protocol): + """The two verification token tables this repository reads and writes.""" + + @property + def litellm_verificationtoken(self) -> TableActions["PrismaVerificationToken"]: ... + + @property + def litellm_deletedverificationtoken(self) -> TableActions["PrismaDeletedVerificationToken"]: ... + + +class _VerificationTokenTransactionManager(Protocol): + async def __aenter__(self) -> _VerificationTokenTables: ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> bool | None: ... + + +class _PrismaVerificationTokenDb(_VerificationTokenTables, Protocol): + def tx(self) -> _VerificationTokenTransactionManager: ... + + +class _PrismaClientView(Protocol): + @property + def db(self) -> _PrismaVerificationTokenDb: ... + _JSON_ENCODED_TOKEN_FIELDS: Final = ( "aliases", @@ -44,17 +74,17 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): """Repository for verification token (API key) database operations.""" @property - def prisma_client(self) -> "PrismaClient": - prisma_client: Final[PrismaClient] = super().prisma_client - return prisma_client + def _db(self) -> _PrismaVerificationTokenDb: + client: Final[_PrismaClientView] = self.prisma_client + return client.db @property def table(self) -> TableActions["PrismaVerificationToken"]: - return self.prisma_client.db.litellm_verificationtoken + return self._db.litellm_verificationtoken @property def deleted_table(self) -> TableActions["PrismaDeletedVerificationToken"]: - return self.prisma_client.db.litellm_deletedverificationtoken + return self._db.litellm_deletedverificationtoken @property def model_class(self) -> type[LiteLLM_VerificationToken]: @@ -325,7 +355,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): archive_data["litellm_changed_by"] = litellm_changed_by archive_data["deleted_at"] = datetime.utcnow() - async with self.prisma_client.db.tx() as tx: + async with self._db.tx() as tx: await tx.litellm_deletedverificationtoken.create(data=archive_data) await tx.litellm_verificationtoken.delete(where={"token": token}) diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index 309894957ea..76e833563ba 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -15,7 +15,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_router_logger if TYPE_CHECKING: - from litellm.types.router import SearchToolTypedDict + from litellm.types.router import SearchToolLiteLLMParams, SearchToolTypedDict class _SearchToolsRouter(Protocol): @@ -34,7 +34,7 @@ class SearchAPIRouter: @staticmethod def _resolve_search_provider_credentials( *, - tool_litellm_params: dict[str, Any], + tool_litellm_params: "SearchToolLiteLLMParams", ) -> tuple[str | None, str | None]: """ Resolve search provider credentials from tool configuration ONLY. From 1ec5083ab4d94c758227c22be052f69e2ae2cf49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:41:07 +0000 Subject: [PATCH 012/246] refactor(types): replace Any with real types across 54 more backend files Third batch of the fifth basedpyright Any reduction round. Every change is typing-only and leaves runtime behavior identical. Provider transformation configs, video and rerank base classes, OTel metadata and the guardrail and realtime type modules move their payload, header and optional-parameter annotations from Any to object, Mapping[str, object] or the concrete model the call site already produces. Repositories and endpoints that reached Prisma through an untyped handle now name the actions they call with the repo's own TableActions protocol. The pydantic field retypes were checked against pydantic to confirm object and Any validate, serialize and generate JSON schema identically. --- .../providers/pydantic_ai_agents/config.py | 4 +-- litellm/integrations/argilla.py | 7 +++-- litellm/integrations/dynamodb.py | 19 +++++++++++-- litellm/integrations/focus/database.py | 6 ++-- .../focus/destinations/vantage_destination.py | 5 ++-- .../langfuse/langfuse_otel_attributes.py | 3 +- litellm/integrations/langsmith.py | 8 +++--- litellm/integrations/otel/model/metadata.py | 14 +++++----- litellm/integrations/otel/mount.py | 19 +++++++++++-- .../litellm_core_utils/coroutine_checker.py | 6 ++-- .../exception_mapping_utils.py | 6 ++-- .../huggingface_template_handler.py | 26 +++++++++++++---- .../anthropic/count_tokens/handler.py | 2 +- .../llms/base_llm/videos/transformation.py | 10 +++---- .../count_tokens/bedrock_token_counter.py | 11 ++++---- litellm/llms/bedrock/files/transformation.py | 6 ++-- .../guardrail_translation/handler.py | 5 +++- litellm/llms/chatgpt/common_utils.py | 4 +-- .../responses/transformation.py | 7 +++-- .../llms/hosted_vllm/chat/transformation.py | 21 ++++++-------- litellm/llms/openai/common_utils.py | 6 ++-- litellm/llms/snowflake/chat/transformation.py | 4 +-- .../stability/image_edit/transformations.py | 4 +-- .../llms/triton/completion/transformation.py | 8 +++--- .../llms/vertex_ai/rerank/transformation.py | 4 +-- .../volcengine/embedding/transformation.py | 9 +++--- litellm/llms/watsonx/rerank/transformation.py | 8 +++--- litellm/llms/xai/realtime/transformation.py | 18 ++++++------ .../analytics_endpoints/cache_activity.py | 20 +++++++++---- .../proxy/client/cli/commands/model_groups.py | 18 ++++++++---- litellm/proxy/client/cli/commands/up.py | 3 +- .../common_utils/cache_pydantic_utils.py | 2 +- .../common_utils/proxy_rate_limit_error.py | 4 +-- .../proxy/container_endpoints/endpoints.py | 10 ++++--- litellm/proxy/db/exception_handler.py | 16 +++++++++-- litellm/proxy/db/spend_log_tool_index.py | 2 +- .../crowdstrike_aidr/crowdstrike_aidr.py | 16 +++++------ .../guardrail_hooks/enkryptai/enkryptai.py | 4 +-- .../guardrail_hooks/qualifire/qualifire.py | 2 +- .../semantic_guard/route_loader.py | 3 +- .../proxy/hooks/parallel_request_limiter.py | 6 ++-- .../common_daily_activity.py | 7 +++-- litellm/proxy/ocr_endpoints/endpoints.py | 14 +++++----- litellm/repositories/budget_repository.py | 28 +++++++++++++++---- .../repositories/organization_repository.py | 28 +++++++++++++++---- litellm/repositories/project_repository.py | 11 ++++---- .../adaptive_router/adaptive_router.py | 9 +++--- .../adaptive_router/signals.py | 8 +++--- .../adaptive_router/update_queue.py | 11 ++++---- litellm/types/containers/main.py | 2 +- litellm/types/guardrails.py | 16 +++++------ litellm/types/integrations/prometheus.py | 8 +++--- litellm/types/realtime.py | 16 +++++------ litellm/vector_store_files/utils.py | 11 ++++---- 54 files changed, 323 insertions(+), 202 deletions(-) diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py index b7546e1a2a1..20404e3702b 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py @@ -23,7 +23,7 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): params: dict[str, Any], api_base: str | None = None, **kwargs: Any, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Handle non-streaming request to Pydantic AI agent.""" if api_base is None: raise ValueError("api_base is required for PydanticAIProviderConfig") @@ -41,7 +41,7 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): params: dict[str, Any], api_base: str | None = None, **kwargs, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """Handle streaming request with fake streaming.""" if not api_base: raise ValueError("api_base is required for Pydantic AI agents") diff --git a/litellm/integrations/argilla.py b/litellm/integrations/argilla.py index 9a87a94cf0b..664ef8efda1 100644 --- a/litellm/integrations/argilla.py +++ b/litellm/integrations/argilla.py @@ -7,7 +7,8 @@ import json import os import random import types -from typing import Any, Final +from collections.abc import Mapping +from typing import Final import httpx from pydantic import BaseModel @@ -69,7 +70,7 @@ class ArgillaLogger(CustomBatchLogger): self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) - def validate_argilla_transformation_object(self, argilla_transformation_object: dict[str, Any]): + def validate_argilla_transformation_object(self, argilla_transformation_object: Mapping[str, object]): if not isinstance(argilla_transformation_object, dict): raise Exception("'argilla_transformation_object' must be a dictionary, to log your payload to Argilla.") @@ -115,7 +116,7 @@ class ArgillaLogger(CustomBatchLogger): ARGILLA_DATASET_NAME=_credentials_dataset_name, ) - def get_chat_messages(self, payload: StandardLoggingPayload) -> list[dict[str, Any]]: + def get_chat_messages(self, payload: StandardLoggingPayload) -> list[dict[str, object]]: payload_messages: Final = payload.get("messages", None) if payload_messages is None: diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py index 38f5924a233..3401ced4efb 100644 --- a/litellm/integrations/dynamodb.py +++ b/litellm/integrations/dynamodb.py @@ -3,12 +3,25 @@ import os import traceback -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, Protocol import litellm from litellm._uuid import uuid +class _DynamoTable(Protocol): + """The one boto3 DynamoDB table call this logger makes.""" + + def put_item(self, *, Item: Mapping[str, object]) -> object: ... + + +class _DynamoResource(Protocol): + """The one boto3 DynamoDB resource call this logger makes.""" + + def Table(self, name: str) -> _DynamoTable: ... + + class DyanmoDBLogger: # Class variables or attributes @@ -16,7 +29,7 @@ class DyanmoDBLogger: # Instance variables import boto3 - self.dynamodb: Any = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"]) + self.dynamodb: Final[_DynamoResource] = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"]) if litellm.dynamodb_table_name is None: raise ValueError( "LiteLLM Error, trying to use DynamoDB but not table name passed. Create a table and set `litellm.dynamodb_table_name=`" @@ -41,7 +54,7 @@ class DyanmoDBLogger: id: Final = response_obj.get("id", str(uuid.uuid4())) # Build the initial payload - payload: Final = { + payload: Final[dict[str, object]] = { "id": id, "call_type": call_type, "startTime": start_time, diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 657c7e0d264..891318f1c54 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import Any, Final +from typing import Final import polars as pl @@ -32,7 +32,7 @@ class FocusLiteLLMDatabase: client: Final = self._ensure_prisma_client() where_clauses: Final[list[str]] = [] - query_params: Final[list[Any]] = [] + query_params: Final[list[datetime | int]] = [] placeholder_index = 1 if start_time_utc: where_clauses.append(f"dus.updated_at >= ${placeholder_index}::timestamptz") @@ -112,7 +112,7 @@ class FocusLiteLLMDatabase: except Exception as exc: raise RuntimeError(f"Error retrieving usage data: {exc}") from exc - async def get_table_info(self) -> dict[str, Any]: + async def get_table_info(self) -> dict[str, object]: """Return metadata about the spend table for diagnostics.""" client: Final = self._ensure_prisma_client() diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index 132f27779c2..68b0d399975 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -4,7 +4,8 @@ from __future__ import annotations import csv import io -from typing import Any, Final +from collections.abc import Mapping +from typing import Final import httpx # noqa: F401 - used at runtime (AsyncClient, HTTPStatusError) @@ -94,7 +95,7 @@ class FocusVantageDestination(FocusDestination): self, *, prefix: str, - config: dict[str, Any] | None = None, + config: Mapping[str, object] | None = None, ) -> None: config = config or {} api_key: Final = config.get("api_key") diff --git a/litellm/integrations/langfuse/langfuse_otel_attributes.py b/litellm/integrations/langfuse/langfuse_otel_attributes.py index 70fea1abb3b..1eb3ce1a9c2 100644 --- a/litellm/integrations/langfuse/langfuse_otel_attributes.py +++ b/litellm/integrations/langfuse/langfuse_otel_attributes.py @@ -5,6 +5,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764 """ import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from pydantic import BaseModel @@ -40,7 +41,7 @@ def get_output_content_by_type( | HttpxBinaryResponseContent | ResponsesAPIResponse | list, - kwargs: dict[str, Any] | None = None, + kwargs: Mapping[str, object] | None = None, ) -> str: """ Extract output content from response objects based on their type. diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 9607eccef52..2c3837b12ec 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -75,9 +75,9 @@ class LangsmithLogger(CustomBatchLogger): if _batch_size: self.batch_size = int(_batch_size) self.log_queue: list[LangsmithQueueObject] = [] - self._flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() + self._flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task() - def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None: + def _start_periodic_flush_task(self) -> asyncio.Task[None] | None: """Start the periodic flush task only when an event loop is already running.""" try: loop: Final = asyncio.get_running_loop() @@ -152,9 +152,9 @@ class LangsmithLogger(CustomBatchLogger): return self._redact_metadata(extra_metadata) - def _build_outputs_with_usage(self, payload: StandardLoggingPayload) -> dict[str, Any]: + def _build_outputs_with_usage(self, payload: StandardLoggingPayload) -> dict[str, object]: response: Final = payload["response"] - outputs: dict[str, Any] + outputs: dict[str, object] if isinstance(response, dict): outputs = {**response} else: diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index ee116aca46b..9eab3021d58 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -36,7 +36,7 @@ model. They coincide on the SDK path, which is correct. from __future__ import annotations -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast @@ -57,7 +57,7 @@ class RequestIdentity: # The team's free-form metadata, carried raw (empty/missing -> None) and # filtered to an operator allowlist only at Baggage-promotion time, so an # unconfigured deployment never promotes any of it. - team_metadata: Mapping[str, Any] | None = None + team_metadata: Mapping[str, object] | None = None key_hash: str | None = None end_user: str | None = None # The model litellm dispatched to the provider. Only known once the call @@ -103,7 +103,7 @@ class RequestIdentity: ``user_api_key_*`` names that ``baggage.DEFAULT_BAGGAGE_METADATA_KEYS`` promotes. """ - get: Final = lambda name: getattr(auth, name, None) # noqa: E731 + get: Final[Callable[[str], object]] = lambda name: getattr(auth, name, None) # noqa: E731 metadata: Final = { meta_key: str(value) for meta_key, attr in ( @@ -217,7 +217,7 @@ class LLMCallEvent: time_to_first_chunk_seconds: float | None @classmethod - def from_dict(cls, kwargs: Mapping[str, Any]) -> LLMCallEvent: + def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent: raw_payload: Final = kwargs.get("standard_logging_object") payload: Final = cast("StandardLoggingPayload", raw_payload) if raw_payload else None operation: Final = resolve_operation(as_str(kwargs.get("call_type"))) @@ -239,7 +239,7 @@ def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None: to the first streamed chunk (``completion_start_time``); ``None`` for non-streaming calls, where ``completion_start_time`` is backfilled with the end time and would not measure first-chunk latency.""" - optional_params: Final = cast(Mapping[str, Any], kwargs.get("optional_params") or {}) + optional_params: Final = cast(Mapping[str, object], kwargs.get("optional_params") or {}) if not optional_params.get("stream"): return None api_call_start: Final = to_seconds(kwargs.get("api_call_start_time")) @@ -307,7 +307,7 @@ def _metadata_dicts( ) -def _call_id(payload: StandardLoggingPayload | None, kwargs: Mapping[str, Any]) -> str | None: +def _call_id(payload: StandardLoggingPayload | None, kwargs: Mapping[str, object]) -> str | None: """The call id from the payload (when closed) or the bare kwargs (at pre_call).""" if payload is not None: call_id: Final = as_str(payload.get("litellm_call_id")) or as_str(payload.get("id")) @@ -351,7 +351,7 @@ def _model_info_id(model_info: object) -> str | None: return None -def _team_metadata_dict(value: object) -> Mapping[str, Any] | None: +def _team_metadata_dict(value: object) -> Mapping[str, object] | None: """The team's free-form metadata as a raw mapping, or ``None`` when missing or empty. diff --git a/litellm/integrations/otel/mount.py b/litellm/integrations/otel/mount.py index ac647c2c4f6..776d8722d14 100644 --- a/litellm/integrations/otel/mount.py +++ b/litellm/integrations/otel/mount.py @@ -12,11 +12,14 @@ when the feature gate is off. """ import os -from typing import Any, Final +from typing import TYPE_CHECKING, Final, Protocol from litellm._logging import verbose_logger from litellm.integrations.otel.model.config import is_otel_v2_enabled +if TYPE_CHECKING: + from fastapi import FastAPI + # Routes excluded from server-span tracing by default: high-frequency pollers and # static UI/docs assets, none of which are LLM traffic. Entries are substring-matched # against the request path (unanchored, so they survive a ``server_root_path`` prefix @@ -65,7 +68,17 @@ PASSTHROUGH_PREFIXES: Final = frozenset( ) -def _passthrough_span_name_hook(span: Any, scope: dict) -> None: +class _RenameableSpan(Protocol): + """The span surface the passthrough naming hook drives.""" + + def is_recording(self) -> bool: ... + + def update_name(self, name: str) -> None: ... + + def set_attribute(self, key: str, value: str) -> None: ... + + +def _passthrough_span_name_hook(span: "_RenameableSpan | None", scope: dict) -> None: """FastAPI ``server_request_hook``: give passthrough server spans a useful name. The instrumentation matches the route at span creation, so both the span name @@ -88,7 +101,7 @@ def _passthrough_span_name_hook(span: Any, scope: dict) -> None: pass -def instrument_fastapi_app(app: Any) -> None: +def instrument_fastapi_app(app: "FastAPI") -> None: """Attach OTel server-span instrumentation to the proxy FastAPI app. Safe no-op when the V2 gate is off or ``opentelemetry-instrumentation-fastapi`` diff --git a/litellm/litellm_core_utils/coroutine_checker.py b/litellm/litellm_core_utils/coroutine_checker.py index 52fc44ba8dc..a5d4f8b968f 100644 --- a/litellm/litellm_core_utils/coroutine_checker.py +++ b/litellm/litellm_core_utils/coroutine_checker.py @@ -16,7 +16,7 @@ class CoroutineChecker: """ def __init__(self): - self._cache = WeakKeyDictionary() + self._cache: WeakKeyDictionary[object, bool] = WeakKeyDictionary() self._max_size = COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY def is_async_callable(self, callback: Any) -> bool: @@ -33,10 +33,10 @@ class CoroutineChecker: pass # Determine target - optimized path for common cases - target = callback + target: object = callback if not inspect.isfunction(target) and not inspect.ismethod(target): try: - call_attr: Final = getattr(target, "__call__", None) + call_attr: Final[object] = getattr(target, "__call__", None) if call_attr is not None: target = call_attr except Exception: diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 82708d412c9..92f07057da5 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,7 +1,7 @@ import json import re import traceback -from typing import Any, Final, Protocol, cast +from typing import Final, Protocol, cast import httpx @@ -191,7 +191,7 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None _response_headers: httpx.Headers | None = None try: _response_headers = getattr(original_exception, "headers", None) - error_response: Final = getattr(original_exception, "response", None) + error_response: Final[object] = getattr(original_exception, "response", None) if not _response_headers and error_response: _response_headers = getattr(error_response, "headers", None) if not _response_headers: @@ -203,7 +203,7 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None def extract_and_raise_litellm_exception( - response: Any | None, + response: object | None, error_str: str, model: str, custom_llm_provider: str, diff --git a/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py b/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py index 3878b36cd91..8f8228d6dfd 100644 --- a/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py +++ b/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py @@ -1,6 +1,8 @@ import json from datetime import datetime -from typing import Any, Final +from typing import Any, Final, Literal + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, @@ -9,6 +11,20 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.types.llms.custom_http import httpxSpecialProvider +class _TokenizerConfigResult(TypedDict): + """Outcome of a tokenizer_config.json fetch, carrying the parsed document when the fetch succeeded.""" + + status: ReadOnly[Literal["success", "failure"]] + tokenizer: NotRequired[ReadOnly[object]] + + +class _ChatTemplateFileResult(TypedDict): + """Outcome of a chat template file fetch, carrying the template body when the fetch succeeded.""" + + status: ReadOnly[Literal["success", "failure"]] + chat_template: NotRequired[ReadOnly[str]] + + def strftime_now(fmt: str) -> str: """ Custom function for templates that need current date/time formatting (e.g., gpt-oss) @@ -22,7 +38,7 @@ def strftime_now(fmt: str) -> str: return datetime.now().strftime(fmt) -def _get_tokenizer_config(hf_model_name: str) -> dict[str, Any]: +def _get_tokenizer_config(hf_model_name: str) -> _TokenizerConfigResult: """ Fetch tokenizer_config.json from HuggingFace (sync) @@ -45,7 +61,7 @@ def _get_tokenizer_config(hf_model_name: str) -> dict[str, Any]: return {"status": "failure"} -async def _aget_tokenizer_config(hf_model_name: str) -> dict[str, Any]: +async def _aget_tokenizer_config(hf_model_name: str) -> _TokenizerConfigResult: """ Fetch tokenizer_config.json from HuggingFace (async) @@ -70,7 +86,7 @@ async def _aget_tokenizer_config(hf_model_name: str) -> dict[str, Any]: return {"status": "failure"} -def _get_chat_template_file(hf_model_name: str) -> dict[str, Any]: +def _get_chat_template_file(hf_model_name: str) -> _ChatTemplateFileResult: """ Fetch chat template from separate .jinja file (sync) @@ -98,7 +114,7 @@ def _get_chat_template_file(hf_model_name: str) -> dict[str, Any]: return {"status": "failure"} -async def _aget_chat_template_file(hf_model_name: str) -> dict[str, Any]: +async def _aget_chat_template_file(hf_model_name: str) -> _ChatTemplateFileResult: """ Fetch chat template from separate .jinja file (async) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 3cc90823af9..36d5a56db0d 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -33,7 +33,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): litellm_params: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + system: object = None, ) -> dict[str, Any]: """ Handle a CountTokens request using httpx with Azure authentication. diff --git a/litellm/llms/base_llm/videos/transformation.py b/litellm/llms/base_llm/videos/transformation.py index f725b295d0f..88cdcd61dd6 100644 --- a/litellm/llms/base_llm/videos/transformation.py +++ b/litellm/llms/base_llm/videos/transformation.py @@ -180,7 +180,7 @@ class BaseVideoConfig(ABC): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video remix request into a URL and data @@ -207,7 +207,7 @@ class BaseVideoConfig(ABC): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video list request into a URL and params @@ -342,8 +342,8 @@ class BaseVideoConfig(ABC): litellm_params: GenericLiteLLMParams, headers: dict, video_file: FileContent | None = None, - extra_body: dict[str, Any] | None = None, - prefetched_source_data: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, + prefetched_source_data: dict[str, object] | None = None, ) -> tuple[str, Mapping[str, object], RequestFiles | None]: """ Transform the video edit request into a URL plus either JSON data or @@ -373,7 +373,7 @@ class BaseVideoConfig(ABC): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video extension request into a URL and JSON data. diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 9c5211ed072..2d02b152c61 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -2,6 +2,7 @@ Bedrock Token Counter implementation using the CountTokens API. """ +from collections.abc import Mapping, Sequence from typing import Any, Final from litellm._logging import verbose_logger @@ -26,12 +27,12 @@ class BedrockTokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, Any]] | None, - contents: list[dict[str, Any]] | None, + messages: Sequence[Mapping[str, object]] | None, + contents: Sequence[Mapping[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: Sequence[Mapping[str, object]] | None = None, + system: object | None = None, ) -> TokenCountResponse | None: """ Count tokens using AWS Bedrock's CountTokens API. @@ -56,7 +57,7 @@ class BedrockTokenCounter(BaseTokenCounter): litellm_params: Final = deployment.get("litellm_params", {}) # Build request data in the format expected by BedrockCountTokensHandler - request_data: Final[dict[str, Any]] = { + request_data: Final[dict[str, object]] = { "model": model_to_use, "messages": messages, } diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 33b27943ad8..f50d63c9f6f 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -247,7 +247,7 @@ def _validate_file_id_against_configured_buckets( return validate_against(configured_bucket_names[-1]) -def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Response) -> int: +def _uploaded_object_size(litellm_params: Mapping[str, object], response_headers: Mapping[str, str]) -> int: """ S3 answers PutObject with an empty body, so the stored object size comes from the signed request recorded by `transform_create_file_request`, not the response headers. @@ -255,7 +255,7 @@ def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Re uploaded_size: Final = litellm_params.get(UPLOAD_CONTENT_LENGTH_PARAM) if isinstance(uploaded_size, int): return uploaded_size - response_content_length: Final = raw_response.headers.get("Content-Length", "0") + response_content_length: Final = response_headers.get("Content-Length", "0") return int(response_content_length) if response_content_length.isdigit() else 0 @@ -1161,7 +1161,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): filename=filename, created_at=int(time.time()), # Current timestamp status="uploaded", - bytes=_uploaded_object_size(litellm_params=litellm_params, raw_response=raw_response), + bytes=_uploaded_object_size(litellm_params=litellm_params, response_headers=raw_response.headers), object="file", ) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index b87f6196e51..eac8afd767c 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -14,6 +14,9 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.pass_through.guardrail_translation.handler import ( + PassThroughEndpointHandler, + ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging @@ -27,7 +30,7 @@ def _is_converse_endpoint(endpoint: str) -> bool: return bool(parts) and parts[-1] in _CONVERSE_ACTIONS -def _generic_passthrough_handler() -> BaseTranslation: +def _generic_passthrough_handler() -> "PassThroughEndpointHandler": """ Fallback for non-Converse Bedrock routes (e.g. invoke). The generic handler scans the full request/response payload so blocking guardrails diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index 35e32e4172f..fe33219f110 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -268,7 +268,7 @@ def _normalize_litellm_params(litellm_params: Any | None) -> dict: return {} -def get_chatgpt_session_id(litellm_params: Any | None) -> str | None: +def get_chatgpt_session_id(litellm_params: object) -> str | None: params: Final = _normalize_litellm_params(litellm_params) for key in ("litellm_session_id", "session_id"): value = params.get(key) @@ -286,5 +286,5 @@ def get_chatgpt_session_id(litellm_params: Any | None) -> str | None: return None -def ensure_chatgpt_session_id(litellm_params: Any | None) -> str: +def ensure_chatgpt_session_id(litellm_params: object) -> str: return get_chatgpt_session_id(litellm_params) or str(uuid4()) diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index 5a4bb798851..8b85b668cba 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -19,6 +19,7 @@ from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfi from litellm.types.llms.openai import ( ResponseInputParam, ResponsesAPIOptionalRequestParams, + ResponsesAPIStreamingResponse, ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders @@ -129,7 +130,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): model: str, parsed_chunk: dict, logging_obj: LiteLLMLoggingObj, - ) -> Any: + ) -> ResponsesAPIStreamingResponse: parsed_chunk = self._normalize_stream_item_id(parsed_chunk) return super().transform_streaming_response( model=model, @@ -262,7 +263,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): # Return the responses endpoint return f"{effective_api_base}/responses" - def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]: + def _handle_reasoning_item(self, item: dict[str, object]) -> dict[str, object]: """ Handle reasoning items for GitHub Copilot, preserving encrypted_content. @@ -280,7 +281,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): # Filter out None values for known problematic fields, # but preserve encrypted_content even if it exists - filtered_item: Final[dict[str, Any]] = {} + filtered_item: Final[dict[str, object]] = {} for k, v in item.items(): # Always include encrypted_content if present (even if None) if k == "encrypted_content": diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 29dc485732f..92bc857e385 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -28,12 +28,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class HostedVLLMChatConfig(OpenAIGPTConfig): - def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, object]]: """ vLLM chat completions currently accepts only OpenAI function tools. Convert custom tools into function tools so request validation does not fail. """ - converted_tools: Final[list[dict[str, Any]]] = [] + converted_tools: Final[list[dict[str, object]]] = [] for idx, tool in enumerate(tools): if not isinstance(tool, dict): converted_tools.append(tool) @@ -63,17 +63,14 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): "required": ["input"], } - function_tool: dict[str, Any] = { - "type": "function", - "function": { - "name": str(tool_name), - "parameters": tool_parameters, - }, + function_definition: dict[str, object] = { + "name": str(tool_name), + "parameters": tool_parameters, } if isinstance(tool_description, str): - function_tool["function"]["description"] = tool_description + function_definition["description"] = tool_description - converted_tools.append(function_tool) + converted_tools.append({"type": "function", "function": function_definition}) return converted_tools @@ -148,7 +145,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): @overload def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: Literal[True] - ) -> Coroutine[Any, Any, list[AllMessageValues]]: ... + ) -> Coroutine[object, object, list[AllMessageValues]]: ... @overload def _transform_messages( @@ -160,7 +157,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: bool = False - ) -> list[AllMessageValues] | Coroutine[Any, Any, list[AllMessageValues]]: + ) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]: """ Support translating: - video files from file_id or file_data to video_url diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 2db6d78a218..e60b6145b44 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -10,7 +10,7 @@ import ssl import time import uuid from collections.abc import AsyncIterator, Iterator, Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Optional +from typing import TYPE_CHECKING, Final, Literal, NamedTuple, Optional from urllib.parse import urlsplit import httpx @@ -87,8 +87,8 @@ class OpenAIError(BaseLLMException): ################################################################### def drop_params_from_unprocessable_entity_error( e: openai.UnprocessableEntityError | httpx.HTTPStatusError, - data: dict[str, Any], -) -> dict[str, Any]: + data: Mapping[str, object], +) -> dict[str, object]: """ Helper function to read OpenAI UnprocessableEntityError and drop the params that raised an error from the error message. diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index f65b0876202..aeff902f655 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -315,7 +315,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): for msg in messages: if isinstance(msg, dict): role = msg.get("role", "") - content: Any = msg.get("content", "") + content: object = msg.get("content", "") msg_cache_control: object = msg.get("cache_control") else: role = getattr(msg, "role", "") @@ -463,7 +463,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): return body - def _transform_tool_choice_to_anthropic(self, tool_choice: Any) -> dict[str, Any]: + def _transform_tool_choice_to_anthropic(self, tool_choice: object) -> Mapping[str, object]: """ Convert tool_choice from OpenAI format to Anthropic format. diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 0b6052ad593..94711d21b50 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -74,7 +74,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): } # Create a copy to not mutate original - convert TypedDict to regular dict - mapped_params: Final[dict[str, Any]] = dict(image_edit_optional_params) + mapped_params: Final[dict[str, object]] = dict(image_edit_optional_params) for k, v in image_edit_optional_params.items(): if k in param_mapping: @@ -182,7 +182,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): # Build Stability request # Populate multipart form-data as separate text fields (data) and files. # Stability expects prompt/output_format/etc. as normal form fields, not file parts. - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "output_format": "png", # Default to PNG } diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index 98a68ba2c36..3c868b3a96f 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -4,7 +4,7 @@ Translates from OpenAI's `/v1/chat/completions` endpoint to Triton's `/generate` import json from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Final, Literal from httpx import Headers, Response @@ -172,7 +172,7 @@ class TritonConfig(BaseConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "TritonResponseIterator": return TritonResponseIterator( streaming_response=streaming_response, sync_stream=sync_stream, @@ -195,14 +195,14 @@ class TritonGenerateConfig(TritonConfig): ) -> dict: inference_params: Final = optional_params.copy() stream: Final = inference_params.pop("stream", False) - data_for_triton: Final[dict[str, Any]] = { + data_for_triton: Final[dict[str, object]] = { "text_input": prompt_factory(model=model, messages=messages), "parameters": { "max_tokens": int(optional_params.get("max_tokens", DEFAULT_MAX_TOKENS_FOR_TRITON)), + **inference_params, }, "stream": bool(stream), } - data_for_triton["parameters"].update(inference_params) return data_for_triton def transform_response( diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 2ec4f2da79b..2b1ed3a29a8 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -5,7 +5,7 @@ Why separate file? Make it easy to see how transformation works """ from collections.abc import Mapping -from typing import Any, Final +from typing import Final import httpx @@ -227,7 +227,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, diff --git a/litellm/llms/volcengine/embedding/transformation.py b/litellm/llms/volcengine/embedding/transformation.py index 091b9dfd334..7c626c66e9d 100644 --- a/litellm/llms/volcengine/embedding/transformation.py +++ b/litellm/llms/volcengine/embedding/transformation.py @@ -3,7 +3,8 @@ Volcengine Embedding Transformation Transforms OpenAI embedding requests to Volcengine format """ -from typing import Any, Final +from collections.abc import Mapping +from typing import Final import httpx @@ -83,11 +84,11 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): def map_openai_params( self, - non_default_params: dict[str, Any], - optional_params: dict[str, Any], + non_default_params: Mapping[str, object], + optional_params: dict[str, object], model: str, drop_params: bool, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Map OpenAI embedding parameters to Volcengine format. diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 293880b188d..0c81d50a5fe 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -5,8 +5,8 @@ Docs - https://cloud.ibm.com/apidocs/watsonx-ai#text-rerank """ import uuid -from collections.abc import Mapping -from typing import Any, Final, cast +from collections.abc import Mapping, Sequence +from typing import Final, cast import httpx @@ -96,7 +96,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: Sequence[str | Mapping[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, @@ -178,7 +178,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): transformed_results: Final = [] for result in _results: - transformed_result: dict[str, Any] = { + transformed_result: dict[str, object] = { "index": result["index"], "relevance_score": result["score"], } diff --git a/litellm/llms/xai/realtime/transformation.py b/litellm/llms/xai/realtime/transformation.py index e9d16daad7c..5efe125ee60 100644 --- a/litellm/llms/xai/realtime/transformation.py +++ b/litellm/llms/xai/realtime/transformation.py @@ -16,7 +16,7 @@ construction time (see ``handler.py``) so all normalization is isolated here and ``RealTimeStreaming`` stays provider-agnostic. """ -from typing import Any, Final +from typing import Final class XAIRealtimeNormalizer: @@ -58,7 +58,7 @@ class XAIRealtimeNormalizer: # Cache content-part objects keyed by (response_id, item_id, content_index) # so that ``response.content_part.done`` events missing ``part`` can be # back-filled from earlier ``content_part.added`` / delta-done events. - self._content_part_by_key: dict[tuple, dict[str, Any]] = {} + self._content_part_by_key: dict[tuple, dict[str, object]] = {} # --------------------------------------------------------------------------- # Public interface consumed by RealTimeStreaming @@ -140,7 +140,7 @@ class XAIRealtimeNormalizer: } self._content_part_by_key[key] = updated - def _resolve_content_part(self, event: dict) -> dict[str, Any]: + def _resolve_content_part(self, event: dict) -> dict[str, object]: part: Final = event.get("part") if isinstance(part, dict): return part @@ -214,7 +214,7 @@ class XAIRealtimeNormalizer: needs_content: Final = event_type in self._EVENTS_NEEDING_CONTENT_INDEX if not needs_output and not needs_content: return event - patch: Final[dict[str, Any]] = {} + patch: Final[dict[str, object]] = {} if needs_output and "output_index" not in event: patch["output_index"] = 0 if needs_content and "content_index" not in event: @@ -228,8 +228,8 @@ class XAIRealtimeNormalizer: # --------------------------------------------------------------------------- @staticmethod - def _default_ga_usage() -> dict[str, Any]: - default_details: Final[dict[str, Any]] = { + def _default_ga_usage() -> dict[str, object]: + default_details: Final[dict[str, int]] = { "cached_tokens": 0, "text_tokens": 0, "audio_tokens": 0, @@ -243,7 +243,7 @@ class XAIRealtimeNormalizer: } @staticmethod - def _normalize_usage(usage: object, *, empty_as_null: bool) -> dict[str, Any] | None: + def _normalize_usage(usage: object, *, empty_as_null: bool) -> dict[str, object] | None: """Coerce a usage object into the full OpenAI GA shape. ``empty_as_null=True`` for ``response.created`` (usage optional). @@ -253,12 +253,12 @@ class XAIRealtimeNormalizer: return None if not usage: return None if empty_as_null else XAIRealtimeNormalizer._default_ga_usage() - default_details: Final[dict[str, Any]] = { + default_details: Final[dict[str, int]] = { "cached_tokens": 0, "text_tokens": 0, "audio_tokens": 0, } - normalized: Final[dict[str, Any]] = { + normalized: Final[dict[str, object]] = { "total_tokens": usage.get("total_tokens", 0), "input_tokens": usage.get("input_tokens", 0), "output_tokens": usage.get("output_tokens", 0), diff --git a/litellm/proxy/analytics_endpoints/cache_activity.py b/litellm/proxy/analytics_endpoints/cache_activity.py index b87b8eac3ef..952ef5bdcfc 100644 --- a/litellm/proxy/analytics_endpoints/cache_activity.py +++ b/litellm/proxy/analytics_endpoints/cache_activity.py @@ -2,16 +2,26 @@ import asyncio import json from collections.abc import Sequence from datetime import datetime -from typing import TYPE_CHECKING, Final +from typing import Final, Protocol from pydantic import BaseModel, TypeAdapter -if TYPE_CHECKING: - from litellm.proxy.utils import PrismaClient - UNKNOWN_CALL_TYPE: Final = "Unknown" +class _SupportsQueryRaw(Protocol): + """The single database operation the cache-activity queries issue.""" + + async def query_raw(self, query: str, *args: object) -> Sequence[object]: ... + + +class _SupportsRawQueryDb(Protocol): + """A prisma client handle, narrowed to the raw-query surface used here.""" + + @property + def db(self) -> _SupportsQueryRaw: ... + + class CacheActivityGroup(BaseModel): call_type: str api_requests: int @@ -143,7 +153,7 @@ def compute_totals(groups: Sequence[CacheActivityGroup]) -> CacheActivityTotals: async def get_cache_activity( - prisma_client: "PrismaClient", + prisma_client: _SupportsRawQueryDb, start_date: datetime, end_date: datetime, key_aliases: Sequence[str], diff --git a/litellm/proxy/client/cli/commands/model_groups.py b/litellm/proxy/client/cli/commands/model_groups.py index c904e5bed49..367c2063b6b 100644 --- a/litellm/proxy/client/cli/commands/model_groups.py +++ b/litellm/proxy/client/cli/commands/model_groups.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final, Literal import click @@ -5,10 +6,17 @@ import rich import rich.table from ... import Client +from ._cli_context import cli_context_values def create_client(ctx: click.Context) -> Client: - return Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + return Client(base_url=context["base_url"], api_key=context["api_key"]) + + +def _rendered_field(group: Mapping[str, object], key: str, default: str) -> str: + """The rendered value of one model group field, or ``default`` when the group omits it.""" + return str(group.get(key, default)) @click.group(name="model-groups") @@ -46,10 +54,10 @@ def list_model_groups(ctx: click.Context, output_format: Literal["table", "json" for group in groups: table.add_row( - str(group.get("model_group", "")), - str(group.get("mode", "chat")), - str(group.get("input_cost_per_token", "")), - str(group.get("output_cost_per_token", "")), + _rendered_field(group, "model_group", ""), + _rendered_field(group, "mode", "chat"), + _rendered_field(group, "input_cost_per_token", ""), + _rendered_field(group, "output_cost_per_token", ""), ) rich.print(table) diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index b7c02866d6f..45d21b0c0b5 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -166,7 +166,8 @@ def up(ctx: click.Context) -> None: is already running (this does not start one for you). Cursor is not supported: it has no equivalent file-based config to patch. """ - base_url: Final = ctx.obj["base_url"] + ctx_obj: Final[CliContextObj] = ctx.obj + base_url: Final = ctx_obj["base_url"] try: _ensure_fresh_login(ctx) diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py index 725c2b61145..3703cf7c916 100644 --- a/litellm/proxy/common_utils/cache_pydantic_utils.py +++ b/litellm/proxy/common_utils/cache_pydantic_utils.py @@ -37,7 +37,7 @@ class CacheCodec: """ @staticmethod - def serialize(value: Any, model_type: type[T] | None = None) -> Any: + def serialize(value: object, model_type: type[T] | None = None) -> object: """ Encode a value for DualCache / Redis (``json.dumps``-safe). diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index c109da6f571..8d3a587a4cd 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -66,7 +66,7 @@ def map_v3_rate_limit_type( return None -def _coerce_message(detail: Any) -> str: +def _coerce_message(detail: object) -> str: """Best-effort, JSON-friendly stringification of an HTTPException-style detail.""" if detail is None: return "" @@ -144,7 +144,7 @@ class ProxyRateLimitError(HTTPException, RateLimitError): def __init__( self, detail: Any, - headers: Mapping[str, Any] | None = None, + headers: Mapping[str, object] | None = None, category: str | RateLimitErrorCategory = RateLimitErrorCategory.LITELLM_RATE_LIMIT, rate_limit_type: str | RateLimitType | None = None, model: str | None = None, diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index e852eb5d6f9..c1407979f29 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -106,7 +106,7 @@ async def create_container( # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - response: Final = await processor.base_process_llm_request( + response: Final[object] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -216,7 +216,7 @@ async def list_containers( or get_custom_llm_provider_from_request_query(request=request) or "openai" ) - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "query_params": query_params, "model": query_params.get("model"), "order": order, @@ -341,7 +341,7 @@ async def retrieve_container( # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + container: Final[object] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -366,6 +366,7 @@ async def retrieve_container( proxy_logging_obj=proxy_logging_obj, version=version, ) + return container @router.delete( @@ -446,7 +447,7 @@ async def delete_container( # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + deleted_container: Final[object] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -471,6 +472,7 @@ async def delete_container( proxy_logging_obj=proxy_logging_obj, version=version, ) + return deleted_container # Register JSON-configured container file endpoints diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index f469587ab8e..de20416ad9c 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,5 +1,5 @@ from collections.abc import Awaitable, Callable, Iterator -from typing import Any, Final, TypeVar +from typing import Final, Protocol, TypeVar from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -407,8 +407,20 @@ def _coerce_timeout(value: object, fallback: float) -> float: _ReadResultT: Final = TypeVar("_ReadResultT") +class _DBReconnectClient(Protocol): + """The one method `call_with_db_reconnect_retry` needs from a Prisma client.""" + + async def attempt_db_reconnect( + self, + *, + reason: str, + timeout_seconds: float | None = None, + lock_timeout_seconds: float | None = None, + ) -> bool: ... + + async def call_with_db_reconnect_retry( - prisma_client: Any, + prisma_client: _DBReconnectClient, coro_factory: Callable[[], Awaitable[_ReadResultT]], *, reason: str, diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 93bbc567430..7a7982481fb 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -148,4 +148,4 @@ async def flush_tool_usage_transactions( except DB_RETRY_SAFE_ERROR_TYPES: if attempt >= n_retry_times: raise - await asyncio.sleep(2**attempt + random.uniform(0, 1)) + await asyncio.sleep(2.0**attempt + random.uniform(0, 1)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index f7f500b1adc..db6fc238b9b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple, Optiona from fastapi import HTTPException from pydantic import BaseModel, ConfigDict, Field, ValidationError -from typing_extensions import Any, override +from typing_extensions import override from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -78,7 +78,7 @@ class _GuardChatCompletionsResult(BaseModel): """Whether or not the prompt triggered a block detection.""" transformed: bool | None = None """Whether or not the original input was transformed.""" - detectors: dict[str, Any] | None = None + detectors: dict[str, object] | None = None """Result of the policy analyzing and input prompt.""" @@ -146,8 +146,8 @@ def _extract_text_from_message(message: _Message) -> str: return "\n".join(part.text for part in content if isinstance(part, _TextContentPart)) -def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Mapping[str, Any] | None: - merged: Final[dict[str, Any]] = {} +def _merge_metadata_bags(request_data: Mapping[str, object]) -> Mapping[str, object] | None: + merged: Final[dict[str, object]] = {} present = False for bag in (request_data.get("metadata"), request_data.get("litellm_metadata")): if isinstance(bag, Mapping): @@ -313,7 +313,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): self._set_streaming_params(streaming_params_from_litellm_params(litellm_params)) async def _call_crowdstrike_aidr_guard( - self, payload: dict[str, Any], hook_name: str + self, payload: dict[str, object], hook_name: str ) -> _GuardChatCompletionsResult: """ Makes the API call to the CrowdStrike AIDR AI Guard endpoint. @@ -423,7 +423,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): return [_extract_text_from_message(msg) for msg in tail] async def _call_or_fail_open( - self, payload: dict[str, Any], hook_name: str, request_data: dict[str, object] + self, payload: dict[str, object], hook_name: str, request_data: dict[str, object] ) -> _GuardChatCompletionsResult: start_time: Final = time.time() try: @@ -506,7 +506,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): event_type = "output" hook_name = "apply_guardrail (response)" - ai_guard_payload: Final[dict[str, Any]] = { + ai_guard_payload: Final[dict[str, object]] = { "guard_input": guard_input.model_dump(mode="json"), "event_type": event_type, } @@ -521,7 +521,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): if user_id: ai_guard_payload["user_id"] = user_id - extra_info: Final[dict[str, str]] = {} + extra_info: Final[dict[str, object]] = {} user_email: Final = metadata.get("user_api_key_user_email") if user_email: extra_info["user_name"] = user_email diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 89afecafb0f..efe959bd186 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -6,7 +6,7 @@ # +-------------------------------------------------------------+ import os -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterable from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional @@ -465,7 +465,7 @@ class EnkryptAIGuardrails(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[ModelResponseStream], request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index d82944c44ed..eceb54681f6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -383,7 +383,7 @@ class QualifireGuardrail(CustomGuardrail): result: Final = response.json() # Extract response info for logging - qualifire_response: Final = { + qualifire_response: Final[dict[str, object]] = { "score": result.get("score"), "status": result.get("status"), } diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py index 8d5923d1302..b2139779925 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py @@ -6,6 +6,7 @@ then builds a SemanticRouter for prompt matching. """ import os +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import yaml @@ -66,7 +67,7 @@ class SemanticGuardRouteLoader: cls, route_templates: list[str] | None, custom_routes_file: str | None, - custom_routes: list[dict[str, Any]] | None, + custom_routes: Sequence[Mapping[str, object]] | None, global_threshold: float = DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD, ) -> list["Route"]: """Build semantic-router Route objects from templates + custom config.""" diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index b313cb64c3f..d41acadc4dd 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -27,7 +27,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache - Span = _Span | Any + Span = _Span InternalUsageCache = _InternalUsageCache else: Span = Any @@ -75,7 +75,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): current: dict | None, request_count_api_key: str, rate_limit_type: Literal["key", "model_per_key", "user", "customer", "team"], - values_to_update_in_cache: list[tuple[Any, Any]], + values_to_update_in_cache: list[tuple[str, object]], ) -> dict: verbose_proxy_logger.info("Current Usage of %s in this minute: %s", rate_limit_type, current) if current is None: @@ -266,7 +266,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): rpm_limit = sys.maxsize values_to_update_in_cache: list[ - tuple[Any, Any] + tuple[str, object] ] = [] # values that need to get updated in cache, will run a batch_set_cache after this function # ------------ diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index ce6a97708ab..6841cadf972 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -17,6 +17,7 @@ from litellm.proxy.spend_tracking.key_metadata_recovery import ( ) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import PrismaClient +from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import DeletedVerificationTokenRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, @@ -1155,8 +1156,10 @@ async def get_daily_activity( include_current_utc_day=include_current_utc_day, ) + spend_table: Final[TableActions[DailySpendRecord]] = getattr(prisma_client.db, table_name) + # Get total count for pagination - total_count: Final[int] = await getattr(prisma_client.db, table_name).count(where=where_conditions) + total_count: Final[int] = await spend_table.count(where=where_conditions) # Fetch paginated results. # ``date`` alone is not a unique sort key -- a busy tenant has many @@ -1168,7 +1171,7 @@ async def get_daily_activity( # total. Adding ``id`` (the row's UUID primary key, present on both # LiteLLM_DailyUserSpend and LiteLLM_DailyTeamSpend) as a tiebreaker # gives every page a stable cursor (#30164). - daily_spend_data: Final[Sequence[DailySpendRecord]] = await getattr(prisma_client.db, table_name).find_many( + daily_spend_data: Final[Sequence[DailySpendRecord]] = await spend_table.find_many( where=where_conditions, order=[ {"date": "desc"}, diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index ebf4d988fdd..1d0108cc2ba 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -2,7 +2,7 @@ import json from collections.abc import Mapping -from typing import Any, Final, cast +from typing import Final, cast import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, UploadFile @@ -48,7 +48,7 @@ def _build_document_from_upload( ) -def _with_request_format(data: Mapping[str, Any], request: Request) -> Mapping[str, Any]: +def _with_request_format(data: Mapping[str, object], request: Request) -> Mapping[str, object]: """ Resolve the requested response format from the body or the `x-req-format` header. @@ -90,7 +90,7 @@ def _native_response(response: object, fastapi_response: Response) -> Response | ) -async def _parse_multipart_form(request: Request) -> dict[str, Any]: +async def _parse_multipart_form(request: Request) -> dict[str, object]: """ Extract OCR data from a multipart form request. @@ -130,7 +130,7 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]: content_type=uploaded_file.content_type, ) - data: Final[dict[str, Any]] = {"document": document} + data: Final[dict[str, object]] = {"document": document} for field_name, field_value in form.items(): if field_name in ("file", "document"): @@ -154,12 +154,12 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]: return data -async def _parse_ocr_request(request: Request) -> Mapping[str, Any]: +async def _parse_ocr_request(request: Request) -> Mapping[str, object]: """Parse an OCR request and apply the `x-req-format` header, if any.""" return _with_request_format(await _parse_ocr_request_body(request), request) -async def _parse_ocr_request_body(request: Request) -> dict[str, Any]: +async def _parse_ocr_request_body(request: Request) -> dict[str, object]: """ Parse an OCR request, supporting both JSON and multipart form data. @@ -320,7 +320,7 @@ async def ocr( # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) - response: Final = await processor.base_process_llm_request( + response: Final[object] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, diff --git a/litellm/repositories/budget_repository.py b/litellm/repositories/budget_repository.py index 62632ffb5f6..205646c8393 100644 --- a/litellm/repositories/budget_repository.py +++ b/litellm/repositories/budget_repository.py @@ -2,7 +2,8 @@ Budget repository for database operations on LiteLLM_BudgetTable. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final, Protocol from litellm.models.budget import LiteLLM_BudgetTable from litellm.repositories.base_repository import BaseRepository @@ -12,12 +13,27 @@ if TYPE_CHECKING: from prisma import models as prisma_models +class _BudgetDb(Protocol): + """The single Prisma table this repository reaches for on ``prisma_client.db``.""" + + @property + def litellm_budgettable(self) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: ... + + +class _PrismaClientView(Protocol): + """The one attribute this repository reads off the untyped Prisma client wrapper.""" + + @property + def db(self) -> _BudgetDb: ... + + class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): """Repository for budget database operations.""" @property def table(self) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: - return self.prisma_client.db.litellm_budgettable + client: Final[_PrismaClientView] = self.prisma_client + return client.db.litellm_budgettable @property def model_class(self) -> type[LiteLLM_BudgetTable]: @@ -34,12 +50,12 @@ class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): max_parallel_requests: int | None = None, tpm_limit: int | None = None, rpm_limit: int | None = None, - model_max_budget: dict[str, Any] | None = None, + model_max_budget: Mapping[str, object] | None = None, budget_duration: str | None = None, allowed_models: list[str] | None = None, ) -> LiteLLM_BudgetTable: """Create a new budget record.""" - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "created_by": created_by, "updated_by": created_by, } @@ -71,12 +87,12 @@ class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): max_parallel_requests: int | None = None, tpm_limit: int | None = None, rpm_limit: int | None = None, - model_max_budget: dict[str, Any] | None = None, + model_max_budget: Mapping[str, object] | None = None, budget_duration: str | None = None, allowed_models: list[str] | None = None, ) -> LiteLLM_BudgetTable | None: """Update an existing budget record.""" - data: Final[dict[str, Any]] = {"updated_by": updated_by} + data: Final[dict[str, object]] = {"updated_by": updated_by} if max_budget is not None: data["max_budget"] = max_budget if soft_budget is not None: diff --git a/litellm/repositories/organization_repository.py b/litellm/repositories/organization_repository.py index 5a9bd3724e0..47eb8f4a609 100644 --- a/litellm/repositories/organization_repository.py +++ b/litellm/repositories/organization_repository.py @@ -2,7 +2,8 @@ Organization repository for database operations on LiteLLM_OrganizationTable. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final, Protocol from litellm.models.organization import LiteLLM_OrganizationTable from litellm.repositories.base_repository import BaseRepository @@ -12,12 +13,27 @@ if TYPE_CHECKING: from prisma import models as prisma_models +class _OrganizationDb(Protocol): + """The single Prisma table this repository reaches for on ``prisma_client.db``.""" + + @property + def litellm_organizationtable(self) -> TableActions["prisma_models.LiteLLM_OrganizationTable"]: ... + + +class _PrismaClientView(Protocol): + """The one attribute this repository reads off the untyped Prisma client wrapper.""" + + @property + def db(self) -> _OrganizationDb: ... + + class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): """Repository for organization database operations.""" @property def table(self) -> TableActions["prisma_models.LiteLLM_OrganizationTable"]: - return self.prisma_client.db.litellm_organizationtable + client: Final[_PrismaClientView] = self.prisma_client + return client.db.litellm_organizationtable @property def model_class(self) -> type[LiteLLM_OrganizationTable]: @@ -39,12 +55,12 @@ class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): budget_id: str, created_by: str, organization_id: str | None = None, - metadata: dict[str, Any] | None = None, + metadata: Mapping[str, object] | None = None, models: list[str] | None = None, object_permission_id: str | None = None, ) -> LiteLLM_OrganizationTable: """Create a new organization.""" - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "organization_alias": organization_alias, "budget_id": budget_id, "created_by": created_by, @@ -67,12 +83,12 @@ class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): updated_by: str, organization_alias: str | None = None, budget_id: str | None = None, - metadata: dict[str, Any] | None = None, + metadata: Mapping[str, object] | None = None, models: list[str] | None = None, object_permission_id: str | None = None, ) -> LiteLLM_OrganizationTable | None: """Update an organization.""" - data: Final[dict[str, Any]] = {"updated_by": updated_by} + data: Final[dict[str, object]] = {"updated_by": updated_by} if organization_alias is not None: data["organization_alias"] = organization_alias if budget_id is not None: diff --git a/litellm/repositories/project_repository.py b/litellm/repositories/project_repository.py index 48e55efd258..905e813f35e 100644 --- a/litellm/repositories/project_repository.py +++ b/litellm/repositories/project_repository.py @@ -2,7 +2,8 @@ Project repository for database operations on LiteLLM_ProjectTable. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final from litellm.models.project import LiteLLM_ProjectTable from litellm.repositories.base_repository import BaseRepository @@ -43,14 +44,14 @@ class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): description: str | None = None, team_id: str | None = None, budget_id: str | None = None, - metadata: dict[str, Any] | None = None, + metadata: Mapping[str, object] | None = None, models: list[str] | None = None, model_rpm_limit: dict[str, int] | None = None, model_tpm_limit: dict[str, int] | None = None, object_permission_id: str | None = None, ) -> LiteLLM_ProjectTable: """Create a new project.""" - data: Final[dict[str, Any]] = { + data: Final[dict[str, object]] = { "created_by": created_by, "updated_by": created_by, } @@ -85,7 +86,7 @@ class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): description: str | None = None, team_id: str | None = None, budget_id: str | None = None, - metadata: dict[str, Any] | None = None, + metadata: Mapping[str, object] | None = None, models: list[str] | None = None, model_rpm_limit: dict[str, int] | None = None, model_tpm_limit: dict[str, int] | None = None, @@ -93,7 +94,7 @@ class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): object_permission_id: str | None = None, ) -> LiteLLM_ProjectTable | None: """Update a project.""" - data: Final[dict[str, Any]] = {"updated_by": updated_by} + data: Final[dict[str, object]] = {"updated_by": updated_by} if project_alias is not None: data["project_alias"] = project_alias if description is not None: diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 12ccacbbc1d..8727c3a69f0 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio import time from collections import OrderedDict +from collections.abc import Mapping from dataclasses import asdict, dataclass from typing import Any, Final, cast @@ -122,7 +123,7 @@ class AdaptiveRouter: prefs = self.model_to_prefs.get(model) or _default_prefs() self._cells[(rt, model)] = initial_cell(prefs, rt) - async def load_state_from_db(self, prisma_client: Any) -> None: + async def load_state_from_db(self, prisma_client: object) -> None: """Add each row's persisted delta to a freshly computed cold-start prior. A row holds an accumulated delta, not a full posterior, and can be one-sided @@ -237,7 +238,7 @@ class AdaptiveRouter: cost_weight=self.config.weights.cost, ) - async def get_state_snapshot(self) -> dict[str, Any]: + async def get_state_snapshot(self) -> dict[str, object]: """In-memory snapshot for the introspection endpoint. Cheap; no DB hit.""" cells: Final = [] for (rt, model), cell in sorted(self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])): @@ -278,7 +279,7 @@ class AdaptiveRouter: @staticmethod def _extract_min_quality_tier( - request_kwargs: dict[str, Any], + request_kwargs: Mapping[str, object], ) -> int | None: """Pull `min_quality_tier` from request headers or metadata. @@ -484,7 +485,7 @@ class AdaptiveRouter: return combined_delta @staticmethod - def _persistable_session_snapshot(state: SessionState) -> dict[str, Any]: + def _persistable_session_snapshot(state: SessionState) -> dict[str, object]: snapshot: Final = asdict(state) for sensitive in ( "last_user_content", diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py index c28613b54eb..72e8d27d2bf 100644 --- a/litellm/router_strategy/adaptive_router/signals.py +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -92,7 +92,7 @@ class Turn: user_content: str | None = None assistant_content: str | None = None - tool_calls: list[dict[str, Any]] = field(default_factory=list) + tool_calls: Sequence[Mapping[str, object]] = field(default_factory=list[Mapping[str, object]]) tool_results: Sequence[Mapping[str, object]] = field(default_factory=list) response_status: int | None = None @@ -174,7 +174,7 @@ def _detect_failure(tool_results: Sequence[Mapping[str, object]]) -> bool: return False -def _signature(call: dict[str, Any]) -> str: +def _signature(call: Mapping[str, Any]) -> str: """Stable signature for loop detection: name + sorted JSON-ish args.""" name: Final = call.get("name") or call.get("function", {}).get("name", "") call_args = call.get("arguments") @@ -185,7 +185,7 @@ def _signature(call: dict[str, Any]) -> str: return f"{name}({call_args})" -def _detect_loop(history: list[str], new_calls: list[dict[str, Any]]) -> bool: +def _detect_loop(history: list[str], new_calls: Sequence[Mapping[str, object]]) -> bool: """Fires if any new call's signature appears >= LOOP_REPEAT_THRESHOLD-1 times in recent history (so this call would be the Nth).""" if not new_calls: @@ -238,7 +238,7 @@ def detect_response_signals( previous_assistant_content: str | None, current_assistant_content: str | None, tool_call_history: list[str], - tool_calls: list[dict[str, Any]], + tool_calls: Sequence[Mapping[str, object]], tool_results: Sequence[Mapping[str, object]], response_status: int | None, ) -> SignalDelta: diff --git a/litellm/router_strategy/adaptive_router/update_queue.py b/litellm/router_strategy/adaptive_router/update_queue.py index 1b9fce284ac..e28f2379f9c 100644 --- a/litellm/router_strategy/adaptive_router/update_queue.py +++ b/litellm/router_strategy/adaptive_router/update_queue.py @@ -19,7 +19,8 @@ to the in-memory aggregator). Flush is async and batched. from __future__ import annotations import asyncio -from typing import Any, Final +from collections.abc import Mapping +from typing import Final from litellm._logging import verbose_router_logger from litellm.repositories.table_repositories import ( @@ -39,7 +40,7 @@ class AdaptiveRouterUpdateQueue: def __init__(self) -> None: self._state_agg: dict[StateKey, dict[str, float]] = {} - self._session_agg: dict[SessionKey, dict[str, Any]] = {} + self._session_agg: dict[SessionKey, Mapping[str, object]] = {} self._lock = asyncio.Lock() self._max_state_size_seen = 0 self._max_session_size_seen = 0 @@ -77,7 +78,7 @@ class AdaptiveRouterUpdateQueue: session_id: str, router_name: str, model_name: str, - state_dict: dict[str, Any], + state_dict: Mapping[str, object], ) -> None: """ Last-write-wins per session row. The state_dict is a snapshot of the @@ -91,7 +92,7 @@ class AdaptiveRouterUpdateQueue: # ---- Flushers (called by background task) ---------------------------- - async def flush_state_to_db(self, prisma_client: Any) -> int: + async def flush_state_to_db(self, prisma_client: object) -> int: """ Drain state aggregator and apply to LiteLLM_AdaptiveRouterState. Returns number of cells flushed. @@ -147,7 +148,7 @@ class AdaptiveRouterUpdateQueue: return len(batch) - async def flush_session_to_db(self, prisma_client: Any) -> int: + async def flush_session_to_db(self, prisma_client: object) -> int: """ Drain session aggregator and upsert into LiteLLM_AdaptiveRouterSession. Returns number of session rows flushed. diff --git a/litellm/types/containers/main.py b/litellm/types/containers/main.py index 6a339fd2eac..62ef524a435 100644 --- a/litellm/types/containers/main.py +++ b/litellm/types/containers/main.py @@ -140,7 +140,7 @@ class ContainerFileObject(BaseModel): created_at: int path: str source: str - _hidden_params: dict[str, Any] = {} + _hidden_params: dict[str, builtins.object] = {} def __contains__(self, key: str) -> bool: return hasattr(self, key) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 02dee40f2a3..175226f8d4b 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -2,10 +2,10 @@ from collections.abc import Mapping from datetime import datetime from enum import Enum from types import MappingProxyType -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( @@ -935,7 +935,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up ), ) - additional_provider_specific_params: dict[str, Any] | None = Field( + additional_provider_specific_params: dict[str, object] | None = Field( default=None, description="Additional provider-specific parameters for generic guardrail APIs", ) @@ -1157,7 +1157,7 @@ class GuardrailEventHooks(str, Enum): class DynamicGuardrailParams(TypedDict): - extra_body: dict[str, Any] + extra_body: ReadOnly[dict[str, object]] class GUARDRAIL_DEFINITION_LOCATION(str, Enum): @@ -1188,7 +1188,7 @@ class GuardrailUIAddGuardrailSettings(BaseModel): supported_modes: list[str] supported_modes_by_provider: dict[str, list[str]] pii_entity_categories: list[PiiEntityCategoryMap] - content_filter_settings: dict[str, Any] | None = None + content_filter_settings: dict[str, object] | None = None class PresidioPerRequestConfig(BaseModel): @@ -1206,8 +1206,8 @@ class ApplyGuardrailRequest(BaseModel): language: str | None = None entities: list[PiiEntityType] | None = None input_type: str = "request" - messages: list[dict[str, Any]] | None = None - metadata: dict[str, Any] | None = None + messages: list[dict[str, object]] | None = None + metadata: dict[str, object] | None = None class ApplyGuardrailResponse(BaseModel): @@ -1217,4 +1217,4 @@ class ApplyGuardrailResponse(BaseModel): class PatchGuardrailRequest(BaseModel): guardrail_name: str | None = None litellm_params: BaseLitellmParams | None = None - guardrail_info: dict[str, Any] | None = None + guardrail_info: dict[str, object] | None = None diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 8498b6f6d00..323f43f4b01 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -44,7 +44,7 @@ def _sanitize_prometheus_label_name(label: str) -> str: _PROMETHEUS_LABEL_VALUE_TRANSLATE_V1: Final = str.maketrans("\n", " ", "\r\u2028\u2029") -def _sanitize_prometheus_label_value(value: Any | None) -> str | None: +def _sanitize_prometheus_label_value(value: object | None) -> str | None: """ Same semantics as :func:`_sanitize_prometheus_label_value`, implemented with ``str.translate`` plus a single escape pass instead of chained ``replace``. @@ -1023,7 +1023,7 @@ class UserAPIKeyLabelValues: ``hashed_api_key``. This supports ``**standard_logging_payload`` in tests. """ field_names: Final = {f.name for f in fields(self)} - merged: Final[dict[str, Any]] = {} + merged: Final[dict[str, object]] = {} for f in fields(self): if f.default_factory is not MISSING: merged[f.name] = f.default_factory() @@ -1060,9 +1060,9 @@ class UserAPIKeyLabelValues: # stays cheap. (Dataclass default `str()` delegates to `__repr__`.) return "" - def model_dump(self) -> dict[str, Any]: + def model_dump(self) -> dict[str, object]: """Same shape as the former Pydantic ``model_dump()`` (plain dict, list tags).""" - d: Final[dict[str, Any]] = {f.name: getattr(self, f.name) for f in fields(self)} + d: Final[dict[str, object]] = {f.name: getattr(self, f.name) for f in fields(self)} d["tags"] = list(self.tags) d["custom_metadata_labels"] = dict(self.custom_metadata_labels) return d diff --git a/litellm/types/realtime.py b/litellm/types/realtime.py index 17dc70126f3..bda3865c46a 100644 --- a/litellm/types/realtime.py +++ b/litellm/types/realtime.py @@ -78,15 +78,15 @@ class RealtimeSessionConfig(BaseModel): type: str | None = None model: str | None = None instructions: str | None = None - audio: dict[str, Any] | None = None + audio: dict[str, object] | None = None include: list[str] | None = None max_output_tokens: int | str | None = None output_modalities: list[str] | None = None - tool_choice: Any | None = None - tools: list[dict[str, Any]] | None = None - tracing: Any | None = None - truncation: Any | None = None - prompt: dict[str, Any] | None = None + tool_choice: object | None = None + tools: list[dict[str, object]] | None = None + tracing: object | None = None + truncation: object | None = None + prompt: dict[str, object] | None = None class RealtimeClientSecretRequest(BaseModel): @@ -114,7 +114,7 @@ class RealtimeClientSecretResponse(BaseModel): expires_at: int | None = None value: str - session: dict[str, Any] | None = None + session: dict[str, object] | None = None class RealtimeTranscriptionSessionRequest(BaseModel): @@ -151,7 +151,7 @@ class RealtimeTranscriptionSessionResponse(BaseModel): model_config = {"extra": "allow"} - client_secret: dict[str, Any] | None = None + client_secret: dict[str, object] | None = None class RealtimeErrorDetail(TypedDict): diff --git a/litellm/vector_store_files/utils.py b/litellm/vector_store_files/utils.py index 94ad5c0ecdf..8b4bff921f8 100644 --- a/litellm/vector_store_files/utils.py +++ b/litellm/vector_store_files/utils.py @@ -1,4 +1,5 @@ -from typing import Any, Final, cast, get_type_hints +from collections.abc import Mapping +from typing import Final, cast, get_type_hints from litellm.types.vector_store_files import ( VectorStoreFileCreateRequest, @@ -11,25 +12,25 @@ class VectorStoreFileRequestUtils: """Helper utilities for constructing vector store file requests.""" @staticmethod - def _filter_params(params: dict[str, Any], model: Any) -> dict[str, Any]: + def _filter_params(params: Mapping[str, object], model: type[object]) -> dict[str, object]: valid_keys: Final = get_type_hints(model).keys() return {key: value for key, value in params.items() if key in valid_keys and value is not None} @staticmethod def get_create_request_params( - params: dict[str, Any], + params: Mapping[str, object], ) -> VectorStoreFileCreateRequest: filtered: Final = VectorStoreFileRequestUtils._filter_params(params=params, model=VectorStoreFileCreateRequest) return cast(VectorStoreFileCreateRequest, filtered) @staticmethod - def get_list_query_params(params: dict[str, Any]) -> VectorStoreFileListQueryParams: + def get_list_query_params(params: Mapping[str, object]) -> VectorStoreFileListQueryParams: filtered = VectorStoreFileRequestUtils._filter_params(params=params, model=VectorStoreFileListQueryParams) return cast(VectorStoreFileListQueryParams, filtered) @staticmethod def get_update_request_params( - params: dict[str, Any], + params: Mapping[str, object], ) -> VectorStoreFileUpdateRequest: filtered: Final = VectorStoreFileRequestUtils._filter_params(params=params, model=VectorStoreFileUpdateRequest) return cast(VectorStoreFileUpdateRequest, filtered) From bf9437780b7fd0aa74827239cfe9ea2e19bbacc7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:09:22 +0000 Subject: [PATCH 013/246] refactor(types): replace Any with real types across 11 more backend files Final batch of the fifth basedpyright Any reduction round. Every change is typing-only and leaves runtime behavior identical. These are the densest remaining files, so the yield per file is small and most of the batch was left alone deliberately. The Riva transcription handler describes the SDK module attributes it reads with Protocols instead of a bare ModuleType, the AWS secret manager stops hiding a botocore header object behind Any, and the sensitive data masker, MCP SSO assertion store and Ovalix guardrail move payload and option annotations to object and Mapping[str, object]. --- .../sensitive_data_masker.py | 10 ++--- .../audio_transcription/handler.py | 37 +++++++++++++--- .../vertex_gemma_models/transformation.py | 7 +-- .../mcp_server/openapi_to_mcp_generator.py | 2 +- .../sso_assertion_store.py | 43 ++++++++++++++++--- litellm/proxy/_lazy_features.py | 4 +- .../guardrail_hooks/ovalix/ovalix.py | 17 +++++--- .../team_callback_endpoints.py | 8 ++-- .../management_endpoints.py | 5 ++- litellm/rag/ingestion/gemini_ingestion.py | 2 +- .../secret_managers/aws_secret_manager_v2.py | 7 ++- 11 files changed, 104 insertions(+), 38 deletions(-) diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 22dd4170963..e78c4bd24c3 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -91,13 +91,13 @@ class SensitiveDataMasker: def _mask_sequence( self, - values: list[Any], + values: Sequence[object], depth: int, max_depth: int, excluded_keys: set[str] | None, key_is_sensitive: bool, - ) -> list[Any]: - masked_items: Final[list[Any]] = [] + ) -> Sequence[object]: + masked_items: Final[list[object]] = [] if depth >= max_depth: return values @@ -197,7 +197,7 @@ def _walk_payload(node: object, key_is_sensitive: bool, depth: int) -> object: return node -def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dict[str, Any]: +def mask_sensitive_keys(data: Mapping[str, object], sensitive_fields: set[str]) -> dict[str, object]: """Return a new dict with values masked for keys listed in ``sensitive_fields``. Unlike :meth:`SensitiveDataMasker.mask_dict`, this does exact key-name @@ -209,7 +209,7 @@ def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dic range and are replaced with a fixed-length all-mask string, so a short credential is never returned verbatim. """ - masked: Final[dict[str, Any]] = {} + masked: Final[dict[str, object]] = {} mask_char: Final = _default_masker.mask_char min_visible: Final = _default_masker.visible_prefix + _default_masker.visible_suffix for key, value in data.items(): diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index d188fac8704..bea77a6761c 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -27,7 +27,6 @@ without the optional STT extras installed. import asyncio import inspect from collections.abc import Callable, Iterable -from types import ModuleType from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.litellm_core_utils.audio_utils.utils import ( @@ -95,11 +94,37 @@ class _AudioEncoding(Protocol): def LINEAR_PCM(self) -> object: ... -def _auth_factory(riva_module: ModuleType) -> Callable[..., _RivaAuth]: +class _RivaClientModule(Protocol): + """The ``riva.client`` entry points this handler calls.""" + + @property + def Auth(self) -> Callable[..., _RivaAuth]: ... + + @property + def ASRService(self) -> Callable[[_RivaAuth], _AsrService]: ... + + +class _RivaAsrModule(Protocol): + """The protobuf constructors this handler calls, from whichever module exposes them.""" + + @property + def AudioEncoding(self) -> _AudioEncoding: ... + + @property + def RecognitionConfig(self) -> Callable[..., _RecognitionConfig]: ... + + @property + def StreamingRecognitionConfig(self) -> Callable[..., _StreamingRecognitionConfig]: ... + + @property + def EndpointingConfig(self) -> Callable[..., _EndpointingConfig]: ... + + +def _auth_factory(riva_module: _RivaClientModule) -> Callable[..., _RivaAuth]: return riva_module.Auth -def _audio_encoding(riva_asr_module: ModuleType) -> _AudioEncoding: +def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding: return riva_asr_module.AudioEncoding @@ -317,7 +342,7 @@ class NvidiaRivaAudioTranscription: def _construct_auth( self, - riva_module: ModuleType, + riva_module: _RivaClientModule, api_base: str, api_key: str | None, optional_params: dict, @@ -349,7 +374,7 @@ class NvidiaRivaAudioTranscription: return _auth_factory(riva_module)(None, use_ssl, api_base, metadata) def _build_recognition_config_proto( - self, riva_asr_module: ModuleType, recognition_config_dict: dict[str, Any] + self, riva_asr_module: _RivaAsrModule, recognition_config_dict: dict[str, Any] ) -> _RecognitionConfig: encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper() encoding_enum: Final[object] = getattr( @@ -436,7 +461,7 @@ class NvidiaRivaAudioTranscription: return final_results -def _import_riva() -> tuple[ModuleType, ModuleType]: +def _import_riva() -> tuple[_RivaClientModule, _RivaAsrModule]: """ Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``. diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 58cf7c7e702..67b01c2dc43 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: import tiktoken from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator class VertexGemmaConfig(OpenAIGPTConfig): @@ -56,7 +57,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): self, model_response: ModelResponse, stream: bool, - ) -> ModelResponse | Any: + ) -> "ModelResponse | MockResponseIterator": """ Helper method to return fake stream iterator if streaming is requested. @@ -138,7 +139,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): client: HTTPHandler | httpx.Client | None, api_base: str, headers: dict[str, str], # mutable-ok: forwarded to post(headers: dict | None) - request_data: dict[str, Any], # mutable-ok: forwarded to post(json: dict | ...) + request_data: dict[str, object], # mutable-ok: forwarded to post(json: dict | ...) timeout: float | httpx.Timeout | None, ) -> httpx.Response: if isinstance(client, HTTPHandler): @@ -173,7 +174,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): client: AsyncHTTPHandler | httpx.AsyncClient | None, api_base: str, headers: dict[str, str], # mutable-ok: forwarded to post(headers: dict | None) - request_data: dict[str, Any], # mutable-ok: forwarded to post(json: dict | ...) + request_data: dict[str, object], # mutable-ok: forwarded to post(json: dict | ...) timeout: float | httpx.Timeout | None, ) -> httpx.Response: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 16f58ef5b76..4cf63b1fceb 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -453,7 +453,7 @@ def _raise_for_upstream_failure( if response.status_code == 401 and relays_upstream_auth: raise MCPUpstreamAuthError( status_code=response.status_code, - www_authenticate=response.headers.get("www-authenticate"), + www_authenticate=dict(response.headers).get("www-authenticate"), server_name=upstream, ) raise MCPOpenApiUpstreamError(response.status_code, upstream) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index f7b92df5ba3..5503d19211b 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -18,6 +18,7 @@ TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against st from __future__ import annotations import json +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from typing import TYPE_CHECKING, Final, Protocol @@ -29,6 +30,8 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS if TYPE_CHECKING: + from prisma.models import LiteLLM_SSOIdentityAssertion + from litellm.proxy.utils import PrismaClient _ASSERTION_DECRYPT_LOG_KEY: Final = "sso_identity_assertion" @@ -36,6 +39,34 @@ _STR_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str) _MAYBE_STR_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) +class _SSOAssertionTable(Protocol): + """The ``LiteLLM_SSOIdentityAssertion`` table operations this store calls.""" + + async def find_unique(self, *, where: Mapping[str, str]) -> LiteLLM_SSOIdentityAssertion | None: ... + + async def find_many(self) -> Sequence[LiteLLM_SSOIdentityAssertion]: ... + + async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> object: ... + + async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> object: ... + + +class _MCPServerTable(Protocol): + """The ``LiteLLM_MCPServerTable`` lookup the retention gate calls.""" + + async def find_first(self, *, where: Mapping[str, str]) -> object | None: ... + + +def _assertion_table(prisma_client: PrismaClient) -> _SSOAssertionTable: + """The SSO assertion table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_ssoidentityassertion + + +def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: + """The MCP server table, typed so the untyped prisma client surface stops here.""" + return prisma_client.db.litellm_mcpservertable + + class SSOIdentityAssertion(BaseModel): """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token, ``expires_at`` bounds its usefulness, and the refresh token renews it without re-login.""" @@ -163,9 +194,7 @@ async def ema_assertion_retention_enabled() -> bool: return True if prisma_client is None: return False - row: Final = await prisma_client.db.litellm_mcpservertable.find_first( - where={"auth_type": MCPAuth.oauth2_id_jag.value} - ) + row: Final = await _mcp_server_table(prisma_client).find_first(where={"auth_type": MCPAuth.oauth2_id_jag.value}) return row is not None @@ -184,7 +213,7 @@ async def persist_sso_identity_assertion( **({"expires_at": assertion.expires_at.isoformat()} if assertion.expires_at else {}), } encoded: Final = _STR_ADAPTER.validate_python(encrypt_value_helper(json.dumps(payload))) - await prisma_client.db.litellm_ssoidentityassertion.upsert( + await _assertion_table(prisma_client).upsert( where={"user_id": user_id}, data={ "create": {"user_id": user_id, "assertion_b64": encoded}, @@ -200,7 +229,7 @@ async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None: if prisma_client is None: return None - row: Final = await prisma_client.db.litellm_ssoidentityassertion.find_unique(where={"user_id": user_id}) + row: Final = await _assertion_table(prisma_client).find_unique(where={"user_id": user_id}) if row is None: return None raw: Final = _MAYBE_STR_ADAPTER.validate_python( @@ -310,13 +339,13 @@ async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, re_encrypted: Final = _STR_ADAPTER.validate_python( encrypt_value_helper(plaintext, new_encryption_key=new_master_key) ) - await prisma_client.db.litellm_ssoidentityassertion.update( + await _assertion_table(prisma_client).update( where={"user_id": row.user_id}, data={"assertion_b64": re_encrypted}, ) return True - rows: Final = await prisma_client.db.litellm_ssoidentityassertion.find_many() + rows: Final = await _assertion_table(prisma_client).find_many() outcomes: Final = [await _rotate_row(row) for row in rows] verbose_proxy_logger.info( "rotate_sso_identity_assertions_master_key: rotated %d row(s), skipped %d", diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 50e0a961a49..17743e153f2 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -13,7 +13,7 @@ from collections.abc import Set as AbstractSet from dataclasses import dataclass, field from typing import TYPE_CHECKING, Final -from starlette.types import Receive, Scope, Send +from starlette.types import ASGIApp, Receive, Scope, Send from litellm._logging import verbose_proxy_logger @@ -266,7 +266,7 @@ class LazyFeatureMiddleware: def __init__( self, - app, + app: ASGIApp, fastapi_app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY_FEATURES, ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index b31ed4b0f4a..c69b24c0553 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -10,6 +10,7 @@ import os from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from typing_extensions import ReadOnly, TypedDict, Unpack from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException @@ -33,6 +34,12 @@ BLOCKED_BY_OVALIX_FALLBACK_MESSAGE: Final = "This message was blocked by Ovalix" BLOCKED_ACTION_TYPE: Final = "block" +class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): + """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" + + supported_event_hooks: ReadOnly[list[GuardrailEventHooks]] + + class OvalixGuardrailMissingSecrets(Exception): """Raised when required Ovalix config (API base, key, application/checkpoint IDs) is missing.""" @@ -80,7 +87,7 @@ class OvalixGuardrail(CustomGuardrail): application_id: str | None = None, pre_checkpoint_id: str | None = None, post_checkpoint_id: str | None = None, - **kwargs: Any, + **kwargs: Unpack[_CustomGuardrailOptions], ): self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE") self._tracker_api_key = tracker_api_key or os.environ.get("OVALIX_TRACKER_API_KEY") @@ -88,10 +95,9 @@ class OvalixGuardrail(CustomGuardrail): self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get("OVALIX_PRE_CHECKPOINT_ID") self._post_checkpoint_id = post_checkpoint_id or os.environ.get("OVALIX_POST_CHECKPOINT_ID") - if "supported_event_hooks" not in kwargs: - kwargs["supported_event_hooks"] = [] + supported_event_hooks: Final = kwargs.get("supported_event_hooks", []) - self._validate_config(kwargs["supported_event_hooks"]) + self._validate_config(supported_event_hooks) self._tracker_headers = httpx.Headers( { @@ -103,7 +109,8 @@ class OvalixGuardrail(CustomGuardrail): self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) - super().__init__(**kwargs) + forwarded: Final[_CustomGuardrailOptions] = {**kwargs, "supported_event_hooks": supported_event_hooks} + super().__init__(**forwarded) verbose_proxy_logger.debug( "Ovalix Guardrail initialized: tracker=%s, application_id=%s, pre_checkpoint_id=%s, post_checkpoint_id=%s", self._tracker_api_base, diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index fe658a13c24..b7fcd4ac7a9 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -473,7 +473,7 @@ async def delete_team_callback( raise _callback_error(404, f"callback_name = {callback_name} is not registered for team_id = {team_id}.") updated_metadata: Final = {**team_metadata, "logging": remaining_callbacks} # mutable-ok: persisted as JSON - encrypted_metadata: Final = encrypt_callback_vars(updated_metadata) + encrypted_metadata: Final[object] = encrypt_callback_vars(updated_metadata) team_metadata_json: Final = json.dumps(encrypted_metadata) updated_team: Final = await TeamRepository(prisma_client).table.update( @@ -610,8 +610,8 @@ async def disable_team_logging( # _get_dynamic_logging_metadata stops at metadata["logging"], where the API # and Admin UI register callbacks, without ever reading callback_settings. team_metadata["logging"] = [] # mutable-ok: the disabled state is persisted as an empty JSON array - team_metadata = encrypt_callback_vars(team_metadata) - team_metadata_json: Final = json.dumps(team_metadata) + encrypted_metadata: Final[object] = encrypt_callback_vars(team_metadata) + team_metadata_json: Final = json.dumps(encrypted_metadata) # Update team in database updated_team: Final = await TeamRepository(prisma_client).table.update( @@ -643,7 +643,7 @@ async def disable_team_logging( await _emit_team_callback_audit_log( team_id=team_id, before_metadata=before_metadata, - after_metadata=team_metadata, + after_metadata=encrypted_metadata, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, ) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 0ca2c4c8865..2fb6813a471 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -88,7 +88,7 @@ def _redact_sensitive_litellm_params(litellm_params: object, _depth: int = 0) -> return None if isinstance(litellm_params, str): try: - parsed: Final = json.loads(litellm_params) + parsed: Final[object] = json.loads(litellm_params) except (TypeError, ValueError): return REDACTED_BY_LITELM_STRING return json.dumps(_redact_sensitive_litellm_params(parsed, _depth + 1)) @@ -589,7 +589,8 @@ async def update_vector_store( try: update_data: Final = data.model_dump(exclude_unset=True) - vector_store_id: Final[str] = update_data.pop("vector_store_id") + vector_store_id: Final[str] = data.vector_store_id + update_data.pop("vector_store_id") # Per-store access control: anyone authenticated who passes the # premium-feature gate could otherwise update *any* vector store — diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index 73a0159fc9f..b81c2cc0ebe 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -277,7 +277,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): raise Exception(error_msg) verbose_logger.debug("Initiate resumable upload response: %s", response.headers) # Extract upload URL from response headers - upload_url: Final = response.headers.get("x-goog-upload-url") + upload_url: Final = dict(response.headers).get("x-goog-upload-url") if not upload_url: raise Exception("No upload URL returned in response headers") diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index e86c8e7c919..9f5bf783958 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -16,7 +16,7 @@ Requires: import json import os -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -35,6 +35,9 @@ from litellm.types.secret_managers.main import KeyManagementSettings from .base_secret_manager import BaseSecretManager +if TYPE_CHECKING: + from botocore.awsrequest import HTTPHeaders + class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): def __init__( @@ -530,7 +533,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): secret_value: str | None = None, optional_params: dict | None = None, request_data: dict | None = None, - ) -> tuple[str, Any, bytes]: + ) -> tuple[str, "HTTPHeaders", bytes]: """Prepare the AWS Secrets Manager request""" try: from botocore.auth import SigV4Auth From a558a0b6a983a78fd24a7e1f3760482c079b6255 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 8 Sep 2026 23:10:03 +0000 Subject: [PATCH 014/246] fix(proxy): renew budget reservation counter TTL while the request is in flight A reservation lives inside spend counter keys that expire on the Redis idle TTL (60s). A stream that outlives the TTL dropped its reservation, so a concurrent request on any worker was admitted against the DB floor until the stream reconciled. Renew the counter TTL with EXPIRE every ttl/2 while the reservation is open and stop once it is finalized, so an idle counter still expires on its own if the worker dies. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 15 +++++ litellm/proxy/proxy_server.py | 6 ++ .../spend_tracking/budget_reservation.py | 44 +++++++++++++- .../proxy/test_budget_reservation.py | 60 ++++++++++++++++++- 4 files changed, 123 insertions(+), 2 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 106c1580110..b1531e299d0 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -76,6 +76,8 @@ class _AsyncRedisCommands(Protocol): def ttl(self, name: str) -> Awaitable[int]: ... + def expire(self, name: str, time: int) -> Awaitable[bool]: ... + def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ... def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... @@ -1795,6 +1797,19 @@ class RedisCache(BaseCache): _record_swallowed_redis_failure(self._circuit_breaker, e) return None + @_redis_circuit_breaker_guard + async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + """EXPIRE an existing key without touching its value. False when the key is absent or Redis failed.""" + _used_ttl: Final = self.get_ttl(ttl=ttl) + if _used_ttl is None: + return False + try: + return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) + except Exception as e: + verbose_logger.debug("Redis EXPIRE Error: %s", e) + _record_swallowed_redis_failure(self._circuit_breaker, e) + return False + @_redis_circuit_breaker_guard async def async_rpush( self, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a617aec9f5c..33f099540b4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3222,6 +3222,12 @@ async def increment_spend_counter(counter_key: str, increment: float): return await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) +async def refresh_spend_counter_ttl(counter_key: str) -> bool: + if spend_counter_cache.redis_cache is None: + return False + return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + + async def _increment_spend_counter_cache(counter_key: str, increment: float): if spend_counter_cache.redis_cache is not None: try: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index d9f6e33c43c..bd9eb154839 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -99,6 +100,42 @@ def get_reserved_counter_keys(budget_reservation: dict | None) -> set: } +_lease_renewals: Final[set[asyncio.Task[None]]] = set() # mutable-ok: asyncio only weak-refs pending tasks + + +def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], counter_keys: frozenset[str]) -> None: + """A reservation lives inside spend counter keys that expire on their Redis TTL. Renew the TTL + while the request is in flight so a request longer than the TTL does not drop its + reservation and admit concurrent requests against the DB floor on any worker.""" + from litellm.proxy.proxy_server import spend_counter_cache + + if spend_counter_cache.redis_cache is None or not counter_keys: + return + task: Final = asyncio.create_task( + _renew_reservation_lease( + budget_reservation=budget_reservation, + counter_keys=counter_keys, + interval=spend_counter_cache.redis_cache.default_ttl / 2, + ) + ) + _lease_renewals.add(task) + task.add_done_callback(_lease_renewals.discard) + + +async def _renew_reservation_lease( + budget_reservation: Mapping[str, object], counter_keys: frozenset[str], interval: float +) -> None: + from litellm.proxy.proxy_server import refresh_spend_counter_ttl + + deadline: Final = time.monotonic() + litellm.request_timeout + while time.monotonic() < deadline: + await asyncio.sleep(interval) + if budget_reservation.get("finalized") is True: + return + for counter_key in counter_keys: + await refresh_spend_counter_ttl(counter_key=counter_key) + + def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: UserAPIKeyAuth | None) -> bool: """ Whether an over-budget key's own ``max_budget`` reservation should be @@ -294,12 +331,17 @@ async def reserve_budget_for_request( llm_router=llm_router, input_token_counts=input_token_counts, ) - return { + budget_reservation: Final = { "reserved_cost": reservation_cost, "entries": applied_entries, "finalized": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), } + _start_reservation_lease_renewal( + budget_reservation=budget_reservation, + counter_keys=frozenset(get_reserved_counter_keys(budget_reservation=budget_reservation)), + ) + return budget_reservation async def reconcile_budget_reservation( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6dab054d8ea..2812be2c53e 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1,5 +1,6 @@ import asyncio import threading +import time from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace @@ -2199,26 +2200,83 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): class _ExpiringRedisCache: - def __init__(self) -> None: + """In-memory stand-in for RedisCache with real wall-clock key expiry.""" + + def __init__(self, default_ttl: float = 60.0) -> None: + self.default_ttl = default_ttl self.store: dict[str, float] = {} + self.expires_at: dict[str, float] = {} + self.refresh_count = 0 + + def _evict_expired(self, key: str) -> None: + if self.expires_at.get(key, float("inf")) <= time.monotonic(): + self.store.pop(key, None) + self.expires_at.pop(key, None) async def async_get_cache(self, key: str, *args: object, **kwargs: object) -> float | None: + self._evict_expired(key) return self.store.get(key) async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self._evict_expired(key) self.store[key] = self.store.get(key, 0.0) + float(value) + self.expires_at[key] = time.monotonic() + self.default_ttl return self.store[key] async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: + self._evict_expired(key) self.store[key] = max(self.store.get(key, float("-inf")), float(value)) + self.expires_at[key] = time.monotonic() + self.default_ttl return self.store[key] async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool: self.store[key] = float(value) + self.expires_at[key] = time.monotonic() + self.default_ttl return True async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: self.store.pop(key, None) + self.expires_at.pop(key, None) + + async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + self._evict_expired(key) + if key not in self.store: + return False + self.refresh_count += 1 + self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl) + return True + + +@pytest.mark.asyncio +async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( + spend_counter_state, +): + """A request that runs longer than the counter TTL must keep its reservation in Redis + (so a concurrent request on any worker still sees it), and renewal must stop once the + reservation is reconciled so an idle counter still expires on its own.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease" + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + + await asyncio.sleep(0.5) + assert await redis_cache.async_get_cache(key=counter_key) == pytest.approx(0.6) + concurrent = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert concurrent is not None + assert concurrent["reserved_cost"] == pytest.approx(0.4) + + await release_budget_reservation(reservation) + await release_budget_reservation(concurrent) + await asyncio.sleep(0.15) + refreshes_after_release = redis_cache.refresh_count + await asyncio.sleep(0.35) + assert redis_cache.refresh_count == refreshes_after_release + assert await redis_cache.async_get_cache(key=counter_key) is None @pytest.mark.asyncio From 3cd4a768aef300c62f66792242904c9bc801c00d Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 8 Sep 2026 23:24:21 +0000 Subject: [PATCH 015/246] fix(proxy): keep the reservation lease alive across a failed Redis EXPIRE Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 9 ++----- litellm/proxy/proxy_server.py | 6 ++++- .../proxy/test_budget_reservation.py | 27 ++++++++++++++++++- 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index b1531e299d0..5ed739b5d10 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1799,16 +1799,11 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: - """EXPIRE an existing key without touching its value. False when the key is absent or Redis failed.""" + """EXPIRE an existing key without touching its value. False when the key is absent.""" _used_ttl: Final = self.get_ttl(ttl=ttl) if _used_ttl is None: return False - try: - return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) - except Exception as e: - verbose_logger.debug("Redis EXPIRE Error: %s", e) - _record_swallowed_redis_failure(self._circuit_breaker, e) - return False + return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) @_redis_circuit_breaker_guard async def async_rpush( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 33f099540b4..476eda3d081 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3225,7 +3225,11 @@ async def increment_spend_counter(counter_key: str, increment: float): async def refresh_spend_counter_ttl(counter_key: str) -> bool: if spend_counter_cache.redis_cache is None: return False - return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + try: + return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + except Exception as e: + verbose_proxy_logger.debug("spend counter TTL refresh skipped for %s: %s", counter_key, e) + return False async def _increment_spend_counter_cache(counter_key: str, increment: float): diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 2812be2c53e..3c33ca11c91 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2202,11 +2202,13 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): class _ExpiringRedisCache: """In-memory stand-in for RedisCache with real wall-clock key expiry.""" - def __init__(self, default_ttl: float = 60.0) -> None: + def __init__(self, default_ttl: float = 60.0, fail_first_refresh: bool = False) -> None: self.default_ttl = default_ttl self.store: dict[str, float] = {} self.expires_at: dict[str, float] = {} + self.refresh_attempts = 0 self.refresh_count = 0 + self.fail_first_refresh = fail_first_refresh def _evict_expired(self, key: str) -> None: if self.expires_at.get(key, float("inf")) <= time.monotonic(): @@ -2239,6 +2241,9 @@ class _ExpiringRedisCache: self.expires_at.pop(key, None) async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + self.refresh_attempts += 1 + if self.fail_first_refresh and self.refresh_attempts == 1: + raise ConnectionError("Redis circuit breaker is open") self._evict_expired(key) if key not in self.store: return False @@ -2279,6 +2284,26 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( assert await redis_cache.async_get_cache(key=counter_key) is None +@pytest.mark.asyncio +async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( + spend_counter_state, +): + """One failed EXPIRE (Redis blip, open circuit breaker) must not end renewal for the + rest of the request.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2, fail_first_refresh=True) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-blip", spend=0.0, max_budget=1.0) + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + + await asyncio.sleep(0.5) + assert redis_cache.refresh_attempts >= 3 + await release_budget_reservation(reservation) + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state, From 668e0142664ac6fa6ab6a6c53925c79cf20b1b6e Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 9 Sep 2026 00:11:05 +0000 Subject: [PATCH 016/246] fix(proxy): stop reservation lease renewal once the request task is gone A streaming /v1/messages client disconnect skips reconciliation, so the lease kept the orphaned reservation alive until request_timeout instead of the plain 60s counter TTL the base branch had. Stop renewing when the request task that took the reservation is done. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/budget_reservation.py | 10 ++++++-- .../proxy/test_budget_reservation.py | 23 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index bd9eb154839..c0af6685600 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -116,6 +116,7 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c budget_reservation=budget_reservation, counter_keys=counter_keys, interval=spend_counter_cache.redis_cache.default_ttl / 2, + request_task=asyncio.current_task(), ) ) _lease_renewals.add(task) @@ -123,14 +124,19 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c async def _renew_reservation_lease( - budget_reservation: Mapping[str, object], counter_keys: frozenset[str], interval: float + budget_reservation: Mapping[str, object], + counter_keys: frozenset[str], + interval: float, + request_task: asyncio.Task[object] | None, ) -> None: + """Stops on finalization or once the request task that took the reservation is gone, so a + disconnect path that skipped reconciliation falls back to the plain counter TTL.""" from litellm.proxy.proxy_server import refresh_spend_counter_ttl deadline: Final = time.monotonic() + litellm.request_timeout while time.monotonic() < deadline: await asyncio.sleep(interval) - if budget_reservation.get("finalized") is True: + if budget_reservation.get("finalized") is True or (request_task is not None and request_task.done()): return for counter_key in counter_keys: await refresh_spend_counter_ttl(counter_key=counter_key) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 3c33ca11c91..626c82ea397 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2304,6 +2304,29 @@ async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_reservation_lease_stops_when_request_task_ends_without_reconciling( + spend_counter_state, +): + """A request whose task ends without reconciling (client disconnect path that skips the + cost callbacks) must not keep renewing: the counter falls back to its plain TTL instead of + pinning the reservation until the request timeout.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-orphan", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease-orphan" + + reservation = await asyncio.create_task(_reserve(valid_token, 0.6, key_cache, proxy_logging_obj)) + assert reservation is not None + assert reservation["finalized"] is False + + await asyncio.sleep(0.5) + assert redis_cache.refresh_count == 0 + assert await redis_cache.async_get_cache(key=counter_key) is None + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state, From 1a358318f8821878ea4623bfd7cb4af2dee3da3f Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 10:38:45 -0700 Subject: [PATCH 017/246] feat(ui): add upgrade banner with latest release changelog stats Adds GET /get/latest_release_info, which fetches the latest stable GitHub release once per hour per worker and buckets its PR bullets by conventional commit prefix into new features, fixes, and other updates. The Admin UI compares the running proxy version to that release and shows a dismissible top banner with the stat line when the proxy is behind. Dismissal is stored per release version in localStorage so the banner returns for the next one. Co-Authored-By: Claude Code --- litellm/proxy/proxy_server.py | 4 + .../latest_release_endpoints.py | 138 +++++++++++++ .../test_latest_release_endpoints.py | 191 ++++++++++++++++++ .../latestRelease/useLatestReleaseInfo.ts | 16 ++ .../src/app/(dashboard)/layout.tsx | 3 + .../src/components/UpgradeBanner.test.tsx | 108 ++++++++++ .../src/components/UpgradeBanner.tsx | 77 +++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 54 +++++ .../src/utils/versionUtils.test.ts | 53 +++++ .../src/utils/versionUtils.ts | 23 +++ 10 files changed, 667 insertions(+) create mode 100644 litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py create mode 100644 tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts create mode 100644 ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx create mode 100644 ui/litellm-dashboard/src/components/UpgradeBanner.tsx create mode 100644 ui/litellm-dashboard/src/utils/versionUtils.test.ts create mode 100644 ui/litellm-dashboard/src/utils/versionUtils.ts diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7219b373dc3..65146e8744f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -647,6 +647,9 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.types_utils.utils import get_instance_fn +from litellm.proxy.ui_crud_endpoints.latest_release_endpoints import ( + router as latest_release_endpoints_router, +) from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( router as ui_crud_endpoints_router, ) @@ -18368,6 +18371,7 @@ app.include_router(debugging_endpoints_router) app.include_router(rust_control_plane_router) app.include_router(ui_crud_endpoints_router) app.include_router(user_banner_endpoints_router) +app.include_router(latest_release_endpoints_router) app.include_router(team_callback_router) app.include_router(budget_management_router) app.include_router(model_management_router) diff --git a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py new file mode 100644 index 00000000000..c0306e65aee --- /dev/null +++ b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py @@ -0,0 +1,138 @@ +import re +from collections import Counter +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, Protocol + +import httpx +from fastapi import APIRouter, Depends +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router: Final = APIRouter() + +LATEST_RELEASE_URL: Final = "https://api.github.com/repos/BerriAI/litellm/releases/latest" +LATEST_RELEASE_FETCH_TIMEOUT_SECONDS: Final = 5 +LATEST_RELEASE_CACHE_TTL_SECONDS: Final = 60 * 60 +LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS: Final = 5 * 60 +LATEST_RELEASE_CACHE_KEY: Final = "latest_release_info" + +_RELEASE_BULLET_PATTERN: Final = re.compile(r"^\*\s+([A-Za-z]+)(\([^)]*\))?!?:\s") + +_Bucket = Literal["new_features", "bug_fixes", "other_updates"] +_PREFIX_BUCKETS: Final[Mapping[str, _Bucket]] = MappingProxyType({"feat": "new_features", "fix": "bug_fixes"}) + + +class LatestReleaseInfo(BaseModel): + version: str + new_features: int + bug_fixes: int + other_updates: int + release_url: str + + +@dataclass(frozen=True, slots=True) +class LatestReleaseUnavailable: + reason: str + + +class _GitHubRelease(BaseModel): + tag_name: str + html_url: str + body: str + + +class _AsyncGetClient(Protocol): + def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ... + + +_latest_release_cache: Final = InMemoryCache(max_size_in_memory=1, default_ttl=LATEST_RELEASE_CACHE_TTL_SECONDS) + + +def _default_client() -> _AsyncGetClient: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + return get_async_httpx_client(llm_provider=httpxSpecialProvider.UI) + + +def _default_cache() -> InMemoryCache: + return _latest_release_cache + + +def count_release_bullets(body: str) -> Counter[_Bucket]: + """ + Bucket a release body's ``* type(scope): title by @user in `` bullets by conventional-commit type. + Lines without that shape (headings, "New Contributors" entries) are skipped, not counted as other. + """ + return Counter( + _PREFIX_BUCKETS.get(match.group(1).lower(), "other_updates") + for line in body.splitlines() + if (match := _RELEASE_BULLET_PATTERN.match(line)) is not None + ) + + +def parse_latest_release(response: httpx.Response) -> LatestReleaseInfo | LatestReleaseUnavailable: + if response.status_code != 200: + return LatestReleaseUnavailable(reason=f"GitHub responded with status {response.status_code}") + try: + release: Final = _GitHubRelease.model_validate_json(response.content) + except ValidationError as e: + return LatestReleaseUnavailable(reason=f"GitHub release payload was not the expected shape: {e}") + counts: Final = count_release_bullets(release.body) + return LatestReleaseInfo( + version=release.tag_name.removeprefix("v"), + new_features=counts["new_features"], + bug_fixes=counts["bug_fixes"], + other_updates=counts["other_updates"], + release_url=release.html_url, + ) + + +async def fetch_latest_release(client: _AsyncGetClient) -> LatestReleaseInfo | LatestReleaseUnavailable: + try: + response: Final = await client.get(LATEST_RELEASE_URL, timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) + except httpx.HTTPError as e: + return LatestReleaseUnavailable(reason=f"{type(e).__name__}: {e}") + return parse_latest_release(response) + + +async def get_latest_release_info( + client: _AsyncGetClient, cache: InMemoryCache +) -> LatestReleaseInfo | LatestReleaseUnavailable: + cached: Final = cache.get_cache(LATEST_RELEASE_CACHE_KEY) + if isinstance(cached, (LatestReleaseInfo, LatestReleaseUnavailable)): + return cached + result: Final = await fetch_latest_release(client) + ttl: Final = ( + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS + if isinstance(result, LatestReleaseUnavailable) + else LATEST_RELEASE_CACHE_TTL_SECONDS + ) + cache.set_cache(LATEST_RELEASE_CACHE_KEY, result, ttl=ttl) + return result + + +@router.get( + "/get/latest_release_info", + tags=["UI Settings"], # mutable-ok: FastAPI's route decorator only accepts a list + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI's route decorator only accepts a list + response_model=LatestReleaseInfo | None, +) +async def latest_release_info( + client: _AsyncGetClient = Depends(_default_client), + cache: InMemoryCache = Depends(_default_cache), +) -> LatestReleaseInfo | None: + """ + Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. + Returns null when GitHub can't be reached so the dashboard upgrade banner simply doesn't render. + """ + result: Final = await get_latest_release_info(client=client, cache=cache) + if isinstance(result, LatestReleaseUnavailable): + verbose_proxy_logger.warning("LiteLLM: latest release info unavailable: %s", result.reason) + return None + return result diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py new file mode 100644 index 00000000000..0bbcdfe317d --- /dev/null +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_latest_release_endpoints.py @@ -0,0 +1,191 @@ +import json +import time +from typing import Final + +import httpx +import pytest +from fastapi.testclient import TestClient + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.proxy_server import app +from litellm.proxy.ui_crud_endpoints.latest_release_endpoints import ( + LATEST_RELEASE_CACHE_KEY, + LATEST_RELEASE_CACHE_TTL_SECONDS, + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS, + LATEST_RELEASE_URL, + LatestReleaseInfo, + LatestReleaseUnavailable, + _default_cache, + _default_client, + count_release_bullets, + get_latest_release_info, +) + +SAMPLE_BODY: Final = """## What's Changed +* feat(proxy): add upgrade banner by @kerry in https://github.com/BerriAI/litellm/pull/1 +* fix(azure): retry on 429 by @a in https://github.com/BerriAI/litellm/pull/2 +* fix: handle empty body by @b in https://github.com/BerriAI/litellm/pull/3 +* Feat(ui)!: drop legacy theme by @c in https://github.com/BerriAI/litellm/pull/4 +* chore(deps): bump httpx by @d in https://github.com/BerriAI/litellm/pull/5 +* docs: fix typo by @e in https://github.com/BerriAI/litellm/pull/6 +* Litellm dev 09 08 2026 by @f in https://github.com/BerriAI/litellm/pull/7 +* refactor(router) : spaced colon does not match by @g in https://github.com/BerriAI/litellm/pull/8 + +## New Contributors +* @kerry made their first contribution in https://github.com/BerriAI/litellm/pull/1 + +**Full Changelog**: https://github.com/BerriAI/litellm/compare/v1.101.0...v1.102.0 +""" + +SAMPLE_RELEASE: Final = { + "tag_name": "v1.102.0", + "html_url": "https://github.com/BerriAI/litellm/releases/tag/v1.102.0", + "body": SAMPLE_BODY, +} +EXPECTED_INFO: Final = { + "version": "1.102.0", + "new_features": 2, + "bug_fixes": 2, + "other_updates": 2, + "release_url": SAMPLE_RELEASE["html_url"], +} + + +class _RecordingClient: + def __init__(self, outcomes: list[httpx.Response | Exception]) -> None: + self._outcomes = outcomes + self.calls: list[tuple[str, float | None]] = [] + + async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + self.calls.append((url, timeout)) + outcome = self._outcomes[min(len(self.calls) - 1, len(self._outcomes) - 1)] + if isinstance(outcome, Exception): + raise outcome + return outcome + + +def _github_response(status: int = 200, payload: object = SAMPLE_RELEASE) -> httpx.Response: + return httpx.Response(status, content=json.dumps(payload).encode()) + + +def _fresh_cache() -> InMemoryCache: + return InMemoryCache(max_size_in_memory=1, default_ttl=LATEST_RELEASE_CACHE_TTL_SECONDS) + + +def _override_dependencies(client: _RecordingClient, cache: InMemoryCache, role: LitellmUserRoles) -> None: + async def auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", user_id="test-user", user_role=role) + + app.dependency_overrides[user_api_key_auth] = auth + app.dependency_overrides[_default_client] = lambda: client + app.dependency_overrides[_default_cache] = lambda: cache + + +@pytest.fixture +def http_client(): + yield TestClient(app) + app.dependency_overrides.pop(user_api_key_auth, None) + app.dependency_overrides.pop(_default_client, None) + app.dependency_overrides.pop(_default_cache, None) + + +class TestCountReleaseBullets: + def test_buckets_by_conventional_commit_type(self): + counts = count_release_bullets(SAMPLE_BODY) + assert counts["new_features"] == 2 + assert counts["bug_fixes"] == 2 + assert counts["other_updates"] == 2 + + def test_ignores_non_bullet_lines_and_contributor_entries(self): + assert ( + sum(count_release_bullets("## What's Changed\n\n* @x made their first contribution in url\n").values()) == 0 + ) + + def test_empty_body_yields_zero_counts(self): + counts = count_release_bullets("") + assert (counts["new_features"], counts["bug_fixes"], counts["other_updates"]) == (0, 0, 0) + + +class TestGetLatestReleaseInfo: + @pytest.mark.asyncio + async def test_fetches_and_parses_github_release(self): + client = _RecordingClient([_github_response()]) + result = await get_latest_release_info(client=client, cache=_fresh_cache()) + assert isinstance(result, LatestReleaseInfo) + assert result.model_dump() == EXPECTED_INFO + assert client.calls == [(LATEST_RELEASE_URL, 5)] + + @pytest.mark.asyncio + async def test_second_call_within_ttl_does_not_refetch(self): + client = _RecordingClient([_github_response()]) + cache = _fresh_cache() + first = await get_latest_release_info(client=client, cache=cache) + second = await get_latest_release_info(client=client, cache=cache) + assert first == second + assert len(client.calls) == 1 + + @pytest.mark.asyncio + async def test_success_is_cached_for_the_full_ttl(self): + cache = _fresh_cache() + await get_latest_release_info(client=_RecordingClient([_github_response()]), cache=cache) + remaining = await cache.async_get_ttl(LATEST_RELEASE_CACHE_KEY) - time.time() + assert LATEST_RELEASE_CACHE_TTL_SECONDS - 5 < remaining <= LATEST_RELEASE_CACHE_TTL_SECONDS + + @pytest.mark.asyncio + async def test_failure_is_cached_briefly_so_github_is_not_hammered(self): + client = _RecordingClient([httpx.ConnectError("boom")]) + cache = _fresh_cache() + first = await get_latest_release_info(client=client, cache=cache) + second = await get_latest_release_info(client=client, cache=cache) + assert isinstance(first, LatestReleaseUnavailable) + assert first == second + assert len(client.calls) == 1 + remaining = await cache.async_get_ttl(LATEST_RELEASE_CACHE_KEY) - time.time() + assert ( + LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS - 5 < remaining <= LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "response", + [ + _github_response(status=403, payload={"message": "rate limited"}), + _github_response(status=500, payload={}), + _github_response(payload={"tag_name": "v1.0.0"}), + httpx.Response(200, content=b"not json"), + ], + ids=["rate_limited", "server_error", "missing_fields", "not_json"], + ) + async def test_bad_github_responses_are_unavailable(self, response: httpx.Response): + result = await get_latest_release_info(client=_RecordingClient([response]), cache=_fresh_cache()) + assert isinstance(result, LatestReleaseUnavailable) + + +class TestLatestReleaseInfoEndpoint: + def test_returns_release_stats_for_authenticated_user(self, http_client): + _override_dependencies(_RecordingClient([_github_response()]), _fresh_cache(), LitellmUserRoles.INTERNAL_USER) + response = http_client.get("/get/latest_release_info") + assert response.status_code == 200 + assert response.json() == EXPECTED_INFO + + def test_returns_null_when_github_is_unreachable(self, http_client): + _override_dependencies( + _RecordingClient([httpx.ConnectError("boom")]), _fresh_cache(), LitellmUserRoles.PROXY_ADMIN + ) + response = http_client.get("/get/latest_release_info") + assert response.status_code == 200 + assert response.json() is None + + def test_repeated_requests_reuse_cache(self, http_client): + client = _RecordingClient([_github_response()]) + _override_dependencies(client, _fresh_cache(), LitellmUserRoles.PROXY_ADMIN) + assert http_client.get("/get/latest_release_info").json() == EXPECTED_INFO + assert http_client.get("/get/latest_release_info").json() == EXPECTED_INFO + assert len(client.calls) == 1 + + def test_rejects_unauthenticated_requests(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-1234") + response = TestClient(app).get("/get/latest_release_info") + assert response.status_code in (401, 403) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts new file mode 100644 index 00000000000..5186baa605c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo.ts @@ -0,0 +1,16 @@ +import { $api } from "@/lib/http/api"; +import type { components } from "@/lib/http/schema"; + +export type LatestReleaseInfo = components["schemas"]["LatestReleaseInfo"]; + +export const useLatestReleaseInfo = (accessToken: string | null | undefined) => + $api.useQuery( + "get", + "/get/latest_release_info", + {}, + { + enabled: Boolean(accessToken), + staleTime: 60 * 60 * 1000, + retry: false, + }, + ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 98f2a36d6f3..928fc5f2410 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -12,6 +12,7 @@ import { DebugWarningBanner } from "@/components/DebugWarningBanner"; import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner"; import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner"; import { UserBanner } from "@/components/UserBanner"; +import { UpgradeBanner } from "@/components/UpgradeBanner"; import { uiHref } from "@/utils/uiHref"; import { PluginModeProvider, usePluginMode } from "@/contexts/PluginModeContext"; import { createApiClient } from "@/lib/http/client"; @@ -115,6 +116,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
@@ -134,6 +136,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { +
{children}
diff --git a/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx b/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx new file mode 100644 index 00000000000..be031ea8ad5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx @@ -0,0 +1,108 @@ +import React from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { fireEvent, render, screen } from "@testing-library/react"; +import { describeRelease, UpgradeBanner, UpgradeBannerView } from "./UpgradeBanner"; +import type { LatestReleaseInfo } from "@/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo"; + +vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails", () => ({ + useHealthReadinessDetails: vi.fn(), +})); +vi.mock("@/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo", () => ({ + useLatestReleaseInfo: vi.fn(), +})); + +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { useLatestReleaseInfo } from "@/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo"; + +const RELEASE: LatestReleaseInfo = { + version: "1.103.0", + new_features: 12, + bug_fixes: 30, + other_updates: 8, + release_url: "https://github.com/BerriAI/litellm/releases/tag/v1.103.0", +}; + +describe("describeRelease", () => { + it("lists features, fixes, and other updates in the agreed order", () => { + expect(describeRelease(RELEASE)).toBe("12 new features, 30 fixes, and 8 other updates"); + }); + + it("singularises counts of one", () => { + expect(describeRelease({ ...RELEASE, new_features: 1, bug_fixes: 1, other_updates: 1 })).toBe( + "1 new feature, 1 fix, and 1 other update", + ); + }); +}); + +describe("UpgradeBannerView", () => { + beforeEach(() => { + localStorage.clear(); + }); + + afterEach(() => { + localStorage.clear(); + }); + + it("renders nothing while either version is unknown", () => { + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + const { container: noRelease } = render(); + expect(noRelease).toBeEmptyDOMElement(); + }); + + it("renders nothing when the running version is up to date or ahead", () => { + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + const { container: ahead } = render(); + expect(ahead).toBeEmptyDOMElement(); + }); + + it("shows the latest version, the stat line, and the current version when behind", () => { + render(); + const alert = screen.getByRole("alert"); + expect(alert).toHaveTextContent("The latest version is v1.103.0: 12 new features, 30 fixes, and 8 other updates"); + expect(alert).toHaveTextContent("Your current version is v1.102.0"); + expect(screen.getByRole("link", { name: "v1.103.0" })).toHaveAttribute("href", RELEASE.release_url); + }); + + it("dismissing hides the banner and keeps it hidden on remount for the same release", () => { + const { unmount } = render(); + fireEvent.click(screen.getByRole("button", { name: "Close" })); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + unmount(); + + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + }); + + it("reappears once a newer release ships after a dismissal", () => { + const { unmount } = render(); + fireEvent.click(screen.getByRole("button", { name: "Close" })); + unmount(); + + render(); + expect(screen.getByRole("alert")).toHaveTextContent("The latest version is v1.104.0"); + }); +}); + +describe("UpgradeBanner", () => { + beforeEach(() => { + localStorage.clear(); + }); + + it("feeds both hooks the access token and renders from their data", () => { + vi.mocked(useHealthReadinessDetails).mockReturnValue({ data: { litellm_version: "1.102.0" } } as any); + vi.mocked(useLatestReleaseInfo).mockReturnValue({ data: RELEASE } as any); + render(); + expect(useHealthReadinessDetails).toHaveBeenCalledWith("token"); + expect(useLatestReleaseInfo).toHaveBeenCalledWith("token"); + expect(screen.getByRole("alert")).toHaveTextContent("The latest version is v1.103.0"); + }); + + it("renders nothing when the release endpoint returns null", () => { + vi.mocked(useHealthReadinessDetails).mockReturnValue({ data: { litellm_version: "1.102.0" } } as any); + vi.mocked(useLatestReleaseInfo).mockReturnValue({ data: null } as any); + const { container } = render(); + expect(container).toBeEmptyDOMElement(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UpgradeBanner.tsx b/ui/litellm-dashboard/src/components/UpgradeBanner.tsx new file mode 100644 index 00000000000..cc57223eaff --- /dev/null +++ b/ui/litellm-dashboard/src/components/UpgradeBanner.tsx @@ -0,0 +1,77 @@ +"use client"; + +import React, { useState } from "react"; +import { ArrowUpCircle, X } from "lucide-react"; +import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { Button } from "@/components/ui/button"; +import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; +import { + type LatestReleaseInfo, + useLatestReleaseInfo, +} from "@/app/(dashboard)/hooks/latestRelease/useLatestReleaseInfo"; +import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; +import { isNewerVersion } from "@/utils/versionUtils"; + +const DISMISS_KEY_PREFIX = "litellm:upgradeBannerDismissed:"; + +interface UpgradeBannerProps { + accessToken: string | null; +} + +interface UpgradeBannerViewProps { + currentVersion: string | null | undefined; + latestRelease: LatestReleaseInfo | null | undefined; +} + +const plural = (count: number, singular: string, pluralForm: string): string => + `${count} ${count === 1 ? singular : pluralForm}`; + +export const describeRelease = ({ new_features, bug_fixes, other_updates }: LatestReleaseInfo): string => + [ + plural(new_features, "new feature", "new features"), + plural(bug_fixes, "fix", "fixes"), + `and ${plural(other_updates, "other update", "other updates")}`, + ].join(", "); + +export const UpgradeBannerView: React.FC = ({ currentVersion, latestRelease }) => { + const [locallyDismissed, setLocallyDismissed] = useState(false); + + if (!currentVersion || !latestRelease || !isNewerVersion(currentVersion, latestRelease.version)) { + return null; + } + + const dismissKey = `${DISMISS_KEY_PREFIX}${latestRelease.version}`; + if (locallyDismissed || getLocalStorageItem(dismissKey) === "true") { + return null; + } + + const handleClose = () => { + setLocalStorageItem(dismissKey, "true"); + setLocallyDismissed(true); + }; + + return ( + + + + The latest version is{" "} + + v{latestRelease.version} + + : {describeRelease(latestRelease)} + + Your current version is v{currentVersion} + + + + + ); +}; + +export const UpgradeBanner: React.FC = ({ accessToken }) => { + const { data: healthData } = useHealthReadinessDetails(accessToken); + const { data: latestRelease } = useLatestReleaseInfo(accessToken); + return ; +}; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 83b0d58f2b2..cd12fd10528 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -5109,6 +5109,27 @@ export interface paths { patch?: never; trace?: never; }; + "/get/latest_release_info": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Latest Release Info + * @description Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. + * Returns null when GitHub can't be reached so the dashboard upgrade banner simply doesn't render. + */ + get: operations["latest_release_info_get_latest_release_info_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/get/mcp_semantic_filter_settings": { parameters: { query?: never; @@ -28352,6 +28373,19 @@ export interface components { } & { [key: string]: unknown; }; + /** LatestReleaseInfo */ + LatestReleaseInfo: { + /** Bug Fixes */ + bug_fixes: number; + /** New Features */ + new_features: number; + /** Other Updates */ + other_updates: number; + /** Release Url */ + release_url: string; + /** Version */ + version: string; + }; /** ListAccessGroupsResponse */ ListAccessGroupsResponse: { /** Access Groups */ @@ -47428,6 +47462,26 @@ export interface operations { }; }; }; + latest_release_info_get_latest_release_info_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LatestReleaseInfo"] | null; + }; + }; + }; + }; get_mcp_semantic_filter_settings_get_mcp_semantic_filter_settings_get: { parameters: { query?: never; diff --git a/ui/litellm-dashboard/src/utils/versionUtils.test.ts b/ui/litellm-dashboard/src/utils/versionUtils.test.ts new file mode 100644 index 00000000000..714c7165d6d --- /dev/null +++ b/ui/litellm-dashboard/src/utils/versionUtils.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it } from "vitest"; +import { isNewerVersion, parseReleaseVersion } from "./versionUtils"; + +describe("parseReleaseVersion", () => { + it("reads the numeric components with or without a leading v", () => { + expect(parseReleaseVersion("1.102.0")).toEqual([1, 102, 0]); + expect(parseReleaseVersion("v1.102.0")).toEqual([1, 102, 0]); + }); + + it("ignores a prerelease suffix", () => { + expect(parseReleaseVersion("1.102.0-dev.3")).toEqual([1, 102, 0]); + expect(parseReleaseVersion("1.102.0.rc1")).toEqual([1, 102, 0]); + }); + + it("returns null for strings that are not a release version", () => { + expect(parseReleaseVersion("")).toBeNull(); + expect(parseReleaseVersion("latest")).toBeNull(); + expect(parseReleaseVersion("1.102")).toBeNull(); + }); +}); + +describe("isNewerVersion", () => { + it("is false when the versions are equal", () => { + expect(isNewerVersion("1.102.0", "1.102.0")).toBe(false); + expect(isNewerVersion("1.102.0", "v1.102.0")).toBe(false); + }); + + it("is true when the latest version is ahead on any component", () => { + expect(isNewerVersion("1.102.0", "1.102.1")).toBe(true); + expect(isNewerVersion("1.102.5", "1.103.0")).toBe(true); + expect(isNewerVersion("1.999.9", "2.0.0")).toBe(true); + }); + + it("is false when the running version is already ahead", () => { + expect(isNewerVersion("1.103.0", "1.102.9")).toBe(false); + expect(isNewerVersion("2.0.0", "1.999.9")).toBe(false); + }); + + it("compares components numerically rather than as strings", () => { + expect(isNewerVersion("1.9.0", "1.10.0")).toBe(true); + expect(isNewerVersion("1.10.0", "1.9.0")).toBe(false); + }); + + it("treats a prerelease of the latest version as not behind", () => { + expect(isNewerVersion("1.102.0-dev.1", "1.102.0")).toBe(false); + expect(isNewerVersion("1.101.0-dev.1", "1.102.0")).toBe(true); + }); + + it("is false when either version cannot be parsed", () => { + expect(isNewerVersion("unknown", "1.102.0")).toBe(false); + expect(isNewerVersion("1.102.0", "")).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/versionUtils.ts b/ui/litellm-dashboard/src/utils/versionUtils.ts new file mode 100644 index 00000000000..62fcb97c0ef --- /dev/null +++ b/ui/litellm-dashboard/src/utils/versionUtils.ts @@ -0,0 +1,23 @@ +const RELEASE_VERSION_PATTERN = /^v?(\d+)\.(\d+)\.(\d+)/; + +export const parseReleaseVersion = (version: string): readonly [number, number, number] | null => { + const match = RELEASE_VERSION_PATTERN.exec(version.trim()); + if (!match) { + return null; + } + return [Number(match[1]), Number(match[2]), Number(match[3])]; +}; + +export const isNewerVersion = (current: string, latest: string): boolean => { + const currentParts = parseReleaseVersion(current); + const latestParts = parseReleaseVersion(latest); + if (!currentParts || !latestParts) { + return false; + } + for (let i = 0; i < 3; i += 1) { + if (latestParts[i] !== currentParts[i]) { + return latestParts[i] > currentParts[i]; + } + } + return false; +}; From 7edc79583139309ef7571e3bb5b77828d9cb535e Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 15:21:39 -0700 Subject: [PATCH 018/246] fix(ui): fix CI failures on upgrade banner PR Use Annotated[..., Depends(...)] instead of a call in the default value to clear the B008 ruff budget. Extract an inline test object over the no-large-inline-object-arg eslint budget. Mock UpgradeBanner in the layout test, matching the other top banners, since it now renders through react-query hooks that need a QueryClientProvider the test doesn't set up. Co-Authored-By: Claude Code --- litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py | 6 +++--- ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx | 4 ++++ ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx | 5 ++--- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py index c0306e65aee..4a01728dfe1 100644 --- a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py @@ -3,7 +3,7 @@ from collections import Counter from collections.abc import Awaitable, Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Final, Literal, Protocol +from typing import Annotated, Final, Literal, Protocol import httpx from fastapi import APIRouter, Depends @@ -124,8 +124,8 @@ async def get_latest_release_info( response_model=LatestReleaseInfo | None, ) async def latest_release_info( - client: _AsyncGetClient = Depends(_default_client), - cache: InMemoryCache = Depends(_default_cache), + client: Annotated[_AsyncGetClient, Depends(_default_client)], + cache: Annotated[InMemoryCache, Depends(_default_cache)], ) -> LatestReleaseInfo | None: """ Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 7f1f4cc4bd5..e1fb6e21596 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -37,6 +37,10 @@ vi.mock("@/components/UserBanner", () => ({ UserBanner: () => null, })); +vi.mock("@/components/UpgradeBanner", () => ({ + UpgradeBanner: () => null, +})); + vi.mock("@/contexts/ThemeContext", () => ({ ThemeProvider: ({ children }: { children: React.ReactNode }) => <>{children}, })); diff --git a/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx b/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx index be031ea8ad5..6aa384600c8 100644 --- a/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/UpgradeBanner.test.tsx @@ -28,9 +28,8 @@ describe("describeRelease", () => { }); it("singularises counts of one", () => { - expect(describeRelease({ ...RELEASE, new_features: 1, bug_fixes: 1, other_updates: 1 })).toBe( - "1 new feature, 1 fix, and 1 other update", - ); + const singularCounts = { ...RELEASE, new_features: 1, bug_fixes: 1, other_updates: 1 }; + expect(describeRelease(singularCounts)).toBe("1 new feature, 1 fix, and 1 other update"); }); }); From daa665e578fcdde3e9a7d1a780581d0e3558b682 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 13 Sep 2026 21:46:25 -0700 Subject: [PATCH 019/246] build(deps): re-suppress GHSA-h7x2-h6g9-p789 in osv-scan, mlflow still has no fixed release (#41036) Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- osv-scanner.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/osv-scanner.toml b/osv-scanner.toml index 3e070fc8cf7..482254d4da6 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -5,5 +5,5 @@ reason = "diskcache has no fixed release published; remove this entry once one e [[IgnoredVulns]] id = "GHSA-h7x2-h6g9-p789" -ignoreUntil = 2026-09-14 -reason = "mlflow has no fixed release published; remove this entry once one exists" +ignoreUntil = 2026-10-14 +reason = "mlflow has no fixed release published (3.16.0, 2026-09-04, and master still store gateway secret api_base unvalidated); remove this entry once one exists" From e46c816ec7b42bdab4ac6cdbb7df10022629e693 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 07:31:41 +0000 Subject: [PATCH 020/246] fix: address review feedback on Any reduction - drop redundant Protocol docstrings in dynamodb and otel mount - widen hosted_vllm custom-tool conversion signature from Any to object --- litellm/integrations/dynamodb.py | 4 ---- litellm/integrations/otel/mount.py | 2 -- litellm/llms/hosted_vllm/chat/transformation.py | 4 ++-- 3 files changed, 2 insertions(+), 8 deletions(-) diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py index 3401ced4efb..3fbbfe91ddf 100644 --- a/litellm/integrations/dynamodb.py +++ b/litellm/integrations/dynamodb.py @@ -11,14 +11,10 @@ from litellm._uuid import uuid class _DynamoTable(Protocol): - """The one boto3 DynamoDB table call this logger makes.""" - def put_item(self, *, Item: Mapping[str, object]) -> object: ... class _DynamoResource(Protocol): - """The one boto3 DynamoDB resource call this logger makes.""" - def Table(self, name: str) -> _DynamoTable: ... diff --git a/litellm/integrations/otel/mount.py b/litellm/integrations/otel/mount.py index 776d8722d14..9340f6e9e15 100644 --- a/litellm/integrations/otel/mount.py +++ b/litellm/integrations/otel/mount.py @@ -69,8 +69,6 @@ PASSTHROUGH_PREFIXES: Final = frozenset( class _RenameableSpan(Protocol): - """The span surface the passthrough naming hook drives.""" - def is_recording(self) -> bool: ... def update_name(self, name: str) -> None: ... diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 92bc857e385..32c60bd01b5 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -4,7 +4,7 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions` import json from collections.abc import Coroutine -from typing import Any, Final, Literal, cast, overload +from typing import Final, Literal, cast, overload from litellm.litellm_core_utils.prompt_templates.common_utils import ( _get_image_mime_type_from_url, @@ -28,7 +28,7 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class HostedVLLMChatConfig(OpenAIGPTConfig): - def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, object]]: + def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, object]]) -> list[dict[str, object]]: """ vLLM chat completions currently accepts only OpenAI function tools. Convert custom tools into function tools so request validation does not fail. From bf8df3ab022b81d54b9d8dcc3d97f8002d1bc3a5 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Thu, 3 Sep 2026 16:18:46 +0200 Subject: [PATCH 021/246] hibp support in password policy --- litellm/proxy/_types.py | 11 ++ litellm/proxy/auth/password_policy.py | 68 +++++++++ .../internal_user_endpoints.py | 9 +- litellm/proxy/proxy_server.py | 3 +- litellm/types/llms/custom_http.py | 1 + .../proxy/auth/test_onboarding.py | 140 +++++++++++++++++- .../proxy/auth/test_password_policy.py | 134 +++++++++++++++++ .../test_internal_user_endpoints.py | 34 +++++ tests/test_litellm/proxy/test__types.py | 16 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 10 files changed, 409 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b315a2beac9..9b38b12ab07 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1808,6 +1808,17 @@ class NewUserRequest(GenerateRequestBase): send_invite_email: bool | None = None sso_user_id: str | None = None organizations: list[str] | None = None + password: str | None = None + + @field_validator("password") + @classmethod + def password_not_supported(cls, value: str | None) -> str | None: + if value is not None: + raise ValueError( + "password cannot be set via /user/new. Users set their own password through an " + "invitation link (POST /invitation/new)." + ) + return value class NewUserResponse(GenerateKeyResponse): diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index ab7a565894a..80c3ebeedca 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -4,13 +4,26 @@ Applied at every path that persists a new or changed password for a DB-backed user (``/user/update``, ``/user/bulk_update``, and the invitation onboarding claim flow), so the strength bar is configured in one place instead of per-endpoint. + +Also screens new passwords against known data breaches via the +haveibeenpwned.com (HIBP) k-anonymity range API: only the first 5 characters +of the password's SHA-1 hash ever leave the proxy, and the check fails open +(allows the password) when HIBP is unreachable. """ +import hashlib from collections.abc import Mapping from dataclasses import dataclass from typing import Final +from litellm._logging import verbose_proxy_logger +from litellm._version import version +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.types.llms.custom_http import httpxSpecialProvider + +HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" +HIBP_TIMEOUT_SECONDS: Final = 5.0 DEFAULT_MIN_LENGTH: Final = 12 MIN_ALLOWED_LENGTH: Final = 8 @@ -90,3 +103,58 @@ def validate_password_policy(password: str, general_settings: Mapping[str, objec param="password", code=400, ) + + +def _hibp_client() -> AsyncHTTPHandler: + return get_async_httpx_client( + llm_provider=httpxSpecialProvider.PasswordBreachCheck, + params={"timeout": HIBP_TIMEOUT_SECONDS}, + ) + + +def _is_suffix_in_range_response(response_body: str, hash_suffix: str) -> bool: + for line in response_body.upper().splitlines(): + entry_suffix, _, count = line.strip().partition(":") + if entry_suffix == hash_suffix: + return int(count.strip() or "0") > 0 + return False + + +async def _is_password_breached(password: str, client: AsyncHTTPHandler) -> bool: + # usedforsecurity=False: SHA-1 is only a lookup key into the HIBP dataset, so no security property rests on it + sha1_hex: Final = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + try: + response: Final = await client.get( + f"{HIBP_RANGE_API_BASE}/{sha1_hex[:5]}", + headers={"Add-Padding": "true", "User-Agent": f"litellm-proxy/{version}"}, + ) + response.raise_for_status() + breached: Final = _is_suffix_in_range_response(response.text, sha1_hex[5:]) + except Exception as e: + verbose_proxy_logger.warning("Breached-password check skipped, HIBP lookup failed: %s", e) + return False + return breached + + +async def validate_password_not_breached( + password: str, + general_settings: Mapping[str, object], + client: AsyncHTTPHandler | None = None, +) -> None: + """Raise ``ProxyException`` (400) if ``password`` appears in a known data breach. + + Fails open: an unreachable or misbehaving HIBP allows the password.""" + check_enabled: Final = general_settings.get("password_policy_check_breached_passwords", True) is not False + if not check_enabled: + return + if not await _is_password_breached(password, client if client is not None else _hibp_client()): + return + raise ProxyException( + message=( + "This password appears in known data breaches and cannot be used. " + "Please choose a different password." + ), + type=ProxyErrorTypes.validation_error, + param="password", + code=400, + ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e3efda507f6..f22bad70817 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -29,7 +29,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object -from litellm.proxy.auth.password_policy import validate_password_policy +from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import ( @@ -162,10 +162,11 @@ def _team_membership_table( return team_membership_table -def _hash_password_in_dict(data: dict, general_settings: Mapping[str, object]) -> None: +async def _hash_password_in_dict(data: dict, general_settings: Mapping[str, object]) -> None: """Validate and hash password field in-place if present.""" if "password" in data and data["password"] is not None: validate_password_policy(data["password"], general_settings) + await validate_password_not_breached(data["password"], general_settings) data["password"] = hash_password(data["password"]) @@ -561,7 +562,7 @@ async def new_user( # generate_key_helper_fn only forwards object_permission_id, so without this the entitlement # the caller sent would be dropped on the floor. data_json = await _set_object_permission(data_json=data_json, prisma_client=prisma_client) - _hash_password_in_dict(data_json, general_settings) + data_json.pop("password", None) # always None: NewUserRequest.password_not_supported rejects any other value teams = data.teams if teams is None: teams = check_if_default_team_set() @@ -1449,7 +1450,7 @@ async def _update_single_user_helper( data_json: Final[dict] = user_request.model_dump(exclude_unset=True) non_default_values = _update_internal_user_params(data_json=data_json, data=user_request) - _hash_password_in_dict(non_default_values, general_settings) + await _hash_password_in_dict(non_default_values, general_settings) existing_user_row: BaseModel | None = None if user_request.user_id: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..d2b486b4410 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -331,7 +331,7 @@ from litellm.proxy.auth.model_checks import ( get_mcp_server_ids, get_team_models, ) -from litellm.proxy.auth.password_policy import validate_password_policy +from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy from litellm.proxy.auth.user_api_key_auth import ( _fetch_global_spend_with_event_coordination, user_api_key_auth, @@ -16369,6 +16369,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): ) validate_password_policy(data.password, general_settings) + await validate_password_not_breached(data.password, general_settings) hashed_pw: Final = hash_password(data.password) current_time = litellm.utils.get_utc_datetime() async with prisma_client.db.tx() as tx: diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index d80d7410aae..793893451df 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -31,6 +31,7 @@ class httpxSpecialProvider(str, Enum): UI = "ui" Sandbox = "sandbox" ModelCostMap = "model_cost_map" + PasswordBreachCheck = "password_breach_check" VerifyTypes = str | bool | ssl.SSLContext diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index 524b655b465..8939baedd50 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -8,15 +8,20 @@ Covers the security behavior of: session key only after the password is written """ +import hashlib from datetime import timedelta from unittest.mock import AsyncMock, MagicMock, patch +import httpx import jwt import pytest +import respx from fastapi import HTTPException import litellm -from litellm.proxy._types import InvitationClaim +from litellm.proxy._types import InvitationClaim, ProxyException + +_POLICY_NO_BREACH_CHECK = {"password_policy_check_breached_passwords": False} # --------------------------------------------------------------------------- # Helpers @@ -386,7 +391,9 @@ async def test_claim_token_rejects_concurrent_reuse_before_password_write(): with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.master_key", "sk-test"), - patch("litellm.proxy.proxy_server.general_settings", {}), + patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + "litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK + ), patch( "litellm.proxy.proxy_server.generate_key_helper_fn", new_callable=AsyncMock, @@ -426,7 +433,9 @@ async def test_claim_token_sets_accepted_at_after_password_written(): with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.master_key", "sk-test"), - patch("litellm.proxy.proxy_server.general_settings", {}), + patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + "litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK + ), patch("litellm.proxy.proxy_server.premium_user", False), patch( "litellm.proxy.proxy_server.generate_key_helper_fn", @@ -483,7 +492,9 @@ async def test_claim_token_rolls_back_invite_when_session_key_mint_fails(): with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.master_key", "sk-test"), - patch("litellm.proxy.proxy_server.general_settings", {}), + patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + "litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK + ), patch( "litellm.proxy.proxy_server.generate_key_helper_fn", new_callable=AsyncMock, @@ -505,3 +516,124 @@ async def test_claim_token_rolls_back_invite_when_session_key_mint_fails(): } assert rollback_kwargs["data"]["accepted_at"] is None assert rollback_kwargs["data"]["is_accepted"] is False + + +# --------------------------------------------------------------------------- +# POST /onboarding/claim_token - password policy +# --------------------------------------------------------------------------- + + +def _hibp_url_for(password: str) -> str: + sha1 = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + return f"https://api.pwnedpasswords.com/range/{sha1[:5]}" + + +def _hibp_suffix_for(password: str) -> str: + return hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()[5:] + + +@pytest.mark.asyncio +async def test_claim_token_rejects_short_password_before_consuming_invite(): + """Default policy requires 12 characters; the invite must stay claimable.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite, _make_user()) + request = _make_claim_request(_make_onboarding_token()) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="Sh0rt!pw", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above + ): + with pytest.raises(ProxyException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.code == "400" + assert "at least 12 characters" in exc_info.value.message + prisma.db.litellm_invitationlink.update_many.assert_not_called() + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@respx.mock +async def test_claim_token_rejects_breached_password_before_consuming_invite(): + """A password found in the HIBP corpus must be rejected and never stored.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + password = "P@ssword123456" + respx.get(_hibp_url_for(password)).mock( + return_value=httpx.Response(200, text=f"{_hibp_suffix_for(password)}:1387") + ) + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite, _make_user()) + request = _make_claim_request(_make_onboarding_token()) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password=password, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above + ): + with pytest.raises(ProxyException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.code == "400" + assert "data breaches" in exc_info.value.message + prisma.db.litellm_invitationlink.update_many.assert_not_called() + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@respx.mock +async def test_claim_token_fails_open_when_hibp_unreachable(): + """An HIBP outage must never block onboarding: the claim proceeds.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + password = "NewP@ssw0rd-2026" + respx.get(_hibp_url_for(password)).mock(side_effect=httpx.ConnectError("no route to host")) + + invite = _make_invite(is_accepted=False) + user = _make_user() + prisma = _make_prisma(invite, user) + request = _make_claim_request(_make_onboarding_token()) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password=password, + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: same as above + patch( # test-quality-ok: same as above + "litellm.proxy.proxy_server.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "sk-generated-key", "user_id": "user-123"}, + ), + patch( # test-quality-ok: same as above + "litellm.proxy.proxy_server.get_custom_url", + return_value="http://localhost:4000/", + ), + patch( # test-quality-ok: same as above + "litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation", + return_value=False, + ), + patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""), # test-quality-ok: same as above + ): + result = await claim_onboarding_link(data=data, request=request) + + assert "token" in result + prisma.db.litellm_usertable.update.assert_called_once() diff --git a/tests/test_litellm/proxy/auth/test_password_policy.py b/tests/test_litellm/proxy/auth/test_password_policy.py index f6e7d443907..edc88d21218 100644 --- a/tests/test_litellm/proxy/auth/test_password_policy.py +++ b/tests/test_litellm/proxy/auth/test_password_policy.py @@ -2,22 +2,54 @@ Tests for the configurable password-strength policy in `litellm.proxy.auth.password_policy`, enforced on every path that persists a new or changed password for a locally-managed user. + +The breach-check (HIBP) tests inject a real AsyncHTTPHandler wrapping an +httpx.MockTransport, so no network is touched and nothing is monkeypatched. """ +import hashlib + +import httpx import pytest +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.password_policy import ( DEFAULT_MIN_LENGTH, MIN_ALLOWED_LENGTH, PasswordPolicy, get_password_policy, + validate_password_not_breached, validate_password_policy, ) STRONG_PASSWORD = "Str0ng!Passw0rd" +def _sha1_upper(password: str) -> str: + return hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + + +def _client_with_transport(handler) -> AsyncHTTPHandler: + http_handler = AsyncHTTPHandler() + http_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + return http_handler + + +def _client_never_called() -> AsyncHTTPHandler: + def handler(request: httpx.Request) -> httpx.Response: + raise AssertionError(f"unexpected HTTP call to {request.url}") + + return _client_with_transport(handler) + + +def _client_returning(body: str, status_code: int = 200) -> AsyncHTTPHandler: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(status_code, text=body) + + return _client_with_transport(handler) + + def test_get_password_policy_defaults_to_pif_baseline(): policy = get_password_policy({}) assert policy == PasswordPolicy( @@ -134,3 +166,105 @@ def test_validate_password_policy_rejects_unicode_letter_as_special_character(): def test_validate_password_policy_accepts_real_special_character_with_unicode_letters(): """Same base password as the rejection test above, plus an actual symbol.""" assert validate_password_policy("Passwörd1234!", {}) is None + + +@pytest.mark.asyncio +async def test_breach_check_skipped_when_disabled(): + result = await validate_password_not_breached( + password="password12345", # breached in reality, but the check is off + general_settings={"password_policy_check_breached_passwords": False}, + client=_client_never_called(), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_rejects_breached_password(): + password = "correct horse battery staple" + sha1 = _sha1_upper(password) + body = f"AAAA000000000000000000000000000000A:0\r\n{sha1[5:]}:42\r\nBBBB000000000000000000000000000000B:7" + + with pytest.raises(ProxyException) as exc_info: + await validate_password_not_breached(password=password, general_settings={}, client=_client_returning(body)) + assert exc_info.value.code == "400" + assert exc_info.value.type == ProxyErrorTypes.validation_error + assert exc_info.value.param == "password" + assert "data breaches" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_only_sha1_prefix_leaves_the_proxy(): + password = "a very secret password" + sha1 = _sha1_upper(password) + captured_requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + result = await validate_password_not_breached( + password=password, general_settings={}, client=_client_with_transport(handler) + ) + assert result is None + + (request,) = captured_requests + assert request.url.path == f"/range/{sha1[:5]}" + assert sha1[5:] not in str(request.url) + assert request.headers["Add-Padding"] == "true" + assert "litellm" in request.headers["User-Agent"] + + +@pytest.mark.asyncio +async def test_ignores_padding_entries_with_zero_count(): + """HIBP padding entries (requested via Add-Padding) carry count 0 and must + not be treated as breaches when they collide with the password's suffix.""" + password = "a padded-away password" + sha1 = _sha1_upper(password) + + result = await validate_password_not_breached( + password=password, general_settings={}, client=_client_returning(f"{sha1[5:]}:0") + ) + assert result is None + + +@pytest.mark.asyncio +async def test_accepts_password_absent_from_breach_corpus(): + result = await validate_password_not_breached( + password="a genuinely novel password", + general_settings={}, + client=_client_returning("0018A45C4D1DEF81644B54AB7F969B88D65:1\r\n00D4F6E8FA6EECAD2A3AA415EEC418D38EC:2"), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_breach_check_fails_open_on_network_error(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("no route to host") + + result = await validate_password_not_breached( + password="password12345", # breached, but HIBP is unreachable + general_settings={}, + client=_client_with_transport(handler), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_breach_check_fails_open_on_http_error_status(): + result = await validate_password_not_breached( + password="password12345", + general_settings={}, + client=_client_returning("service unavailable", status_code=503), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_breach_check_fails_open_on_malformed_response_body(): + result = await validate_password_not_breached( + password="password12345", + general_settings={}, + client=_client_returning(f"{_sha1_upper('password12345')[5:]}:not-a-number"), + ) + assert result is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 0d8b19345f1..e7b5172e0fb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1,9 +1,12 @@ +import hashlib import json from datetime import datetime, timezone from types import SimpleNamespace from typing import Final +import httpx import pytest +import respx from fastapi.testclient import TestClient from fastapi import HTTPException from pytest_mock import MockerFixture @@ -4502,6 +4505,11 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo _update_single_user_helper, ) + mocker.patch( # test-quality-ok: same module-global mocking every test in this file already uses + "litellm.proxy.proxy_server.general_settings", + {"password_policy_check_breached_passwords": False}, + ) + mock_prisma_client = _admin_prisma existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user"} @@ -4519,3 +4527,29 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo written_data = mock_prisma_client.update_data.call_args.kwargs["data"] assert written_data.get("password") is not None assert written_data["password"] != strong_password + + +@pytest.mark.asyncio +@respx.mock +async def test_user_update_rejects_breached_password(_admin_prisma): + """A strength-passing password found in the HIBP corpus must be rejected + before it ever reaches the DB write.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + password = "Str0ng!Passw0rd" + sha1 = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + respx.get(f"https://api.pwnedpasswords.com/range/{sha1[:5]}").mock( + return_value=httpx.Response(200, text=f"{sha1[5:]}:1387") + ) + + user_request = UpdateUserRequest(user_id="target-user", password=password) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(ProxyException) as exc_info: + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) + + assert exc_info.value.code == "400" + assert "data breaches" in exc_info.value.message + _admin_prisma.db.litellm_usertable.find_first.assert_not_called() diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index 26bb1533da4..8f9c44d7a38 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -10,6 +10,7 @@ from litellm.proxy._types import ( LiteLLM_AuditLogs, LiteLLM_TeamMembership, LitellmUserRoles, + NewUserRequest, OrganizationMemberUpdateRequest, ResetSpendRequest, UpdateKeyRequest, @@ -277,3 +278,18 @@ def test_team_membership_budget_table_present_still_works(): } result = LiteLLM_TeamMembership.model_validate(data) assert result.litellm_budget_table is None + + +def test_new_user_request_loudly_rejects_a_password(): + """ + /user/new has never persisted a password (the field used to be silently + dropped). Sending one must now fail visibly so the dead path cannot be + revived without going through the password policy. + """ + with pytest.raises(ValidationError, match="invitation link"): + NewUserRequest(user_email="alice@example.com", password="hunter2hunter2") + + +def test_new_user_request_without_password_still_works(): + request = NewUserRequest(user_email="alice@example.com") + assert request.password is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..0cc49382498 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32869,6 +32869,8 @@ export interface components { object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null; /** Organizations */ organizations?: string[] | null; + /** Password */ + password?: string | null; /** * Permissions * @default {} From e2ea7e97a58300c2d885c1f1c04854799fbbfcf6 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Thu, 3 Sep 2026 16:54:26 +0200 Subject: [PATCH 022/246] fix(auth): document /user/new password rejection and format password_policy --- litellm/proxy/auth/password_policy.py | 3 +-- litellm/proxy/management_endpoints/internal_user_endpoints.py | 1 + 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index 80c3ebeedca..84680fa9020 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -151,8 +151,7 @@ async def validate_password_not_breached( return raise ProxyException( message=( - "This password appears in known data breaches and cannot be used. " - "Please choose a different password." + "This password appears in known data breaches and cannot be used. Please choose a different password." ), type=ProxyErrorTypes.validation_error, param="password", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f22bad70817..409877bdc6f 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -495,6 +495,7 @@ async def new_user( - prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts. - organizations: List[str] - List of organization id's the user is a member of - budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. + - password: Optional[str] - Not supported; any value is rejected with a 422. Users set their own password through an invitation link (POST /invitation/new). Returns: - key: (str) The generated api key for the user - expires: (datetime) Datetime object for when key expires. From 1f0ab3d176af4955a1c75b1b44371a310529e169 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Thu, 3 Sep 2026 16:57:28 +0200 Subject: [PATCH 023/246] fix(auth): drop general_settings import left unused in new_user --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 409877bdc6f..ddef2127945 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -515,7 +515,7 @@ async def new_user( ``` """ try: - from litellm.proxy.proxy_server import _license_check, general_settings, prisma_client + from litellm.proxy.proxy_server import _license_check, prisma_client if prisma_client is None: raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value) From fcf7cb6e0c7fc5d39c7f4017630c48becb086ca9 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Thu, 3 Sep 2026 17:12:27 +0200 Subject: [PATCH 024/246] fix(ui): regenerate schema.d.ts for the new_user password docstring --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0cc49382498..c5ee810069d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16798,6 +16798,7 @@ export interface paths { * - prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts. * - organizations: List[str] - List of organization id's the user is a member of * - budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. + * - password: Optional[str] - Not supported; any value is rejected with a 422. Users set their own password through an invitation link (POST /invitation/new). * Returns: * - key: (str) The generated api key for the user * - expires: (datetime) Datetime object for when key expires. From a9a0bcb9f84909cb41e3670b881292d2e415c1b8 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Fri, 4 Sep 2026 21:13:14 +0200 Subject: [PATCH 025/246] move hibp url to constants --- litellm/constants.py | 3 +++ litellm/proxy/auth/password_policy.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..59a7c3d1e40 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2039,3 +2039,6 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) + +# API endpoint for breached password k-anonymity search +HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" \ No newline at end of file diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index 84680fa9020..958547d6fc7 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -18,11 +18,11 @@ from typing import Final from litellm._logging import verbose_proxy_logger from litellm._version import version +from litellm.constants import HIBP_RANGE_API_BASE from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.types.llms.custom_http import httpxSpecialProvider -HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" HIBP_TIMEOUT_SECONDS: Final = 5.0 DEFAULT_MIN_LENGTH: Final = 12 From 0bb0218d0b60dddb1a05ad2e5ef498e417e210e9 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Tue, 8 Sep 2026 16:27:06 +0200 Subject: [PATCH 026/246] fix(auth): screen bulk-update passwords concurrently before any db write /user/bulk_update awaited a separate HIBP lookup for each user in the batch, so a degraded-slow HIBP (5s timeout per lookup) could stretch a 500-user batch to ~2500s and time out the request after some updates had already persisted. validate_passwords_bulk dedupes the batch's passwords, strength-checks first, then fires every needed HIBP lookup concurrently, bounding the worst case at one 5s timeout window. bulk_update_processed_users now screens the whole batch before the serial update loop, so a rejected password fails only its own entry and validation failures precede any persistence. --- litellm/constants.py | 2 +- litellm/proxy/auth/password_policy.py | 79 ++- .../internal_user_endpoints.py | 43 +- .../proxy/auth/test_password_policy.py | 75 +++ .../test_internal_user_endpoints.py | 575 +++++++----------- 5 files changed, 390 insertions(+), 384 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 59a7c3d1e40..5f8fa203b37 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2041,4 +2041,4 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" EMPTY_MAPPING: Final = MappingProxyType({}) # API endpoint for breached password k-anonymity search -HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" \ No newline at end of file +HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range" diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index 958547d6fc7..c27f276252a 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -11,9 +11,11 @@ of the password's SHA-1 hash ever leave the proxy, and the check fails open (allows the password) when HIBP is unreachable. """ +import asyncio import hashlib -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import Final from litellm._logging import verbose_proxy_logger @@ -136,6 +138,33 @@ async def _is_password_breached(password: str, client: AsyncHTTPHandler) -> bool return breached +def is_breach_check_enabled(general_settings: Mapping[str, object]) -> bool: + return general_settings.get("password_policy_check_breached_passwords", True) is not False + + +async def is_password_breached( + password: str, + general_settings: Mapping[str, object], + client: AsyncHTTPHandler | None = None, +) -> bool: + """False when the check is disabled, the password is absent from the HIBP + corpus, or HIBP is unreachable (fail open).""" + if not is_breach_check_enabled(general_settings): + return False + return await _is_password_breached(password, client if client is not None else _hibp_client()) + + +def breached_password_error() -> ProxyException: + return ProxyException( + message=( + "This password appears in known data breaches and cannot be used. Please choose a different password." + ), + type=ProxyErrorTypes.validation_error, + param="password", + code=400, + ) + + async def validate_password_not_breached( password: str, general_settings: Mapping[str, object], @@ -144,16 +173,42 @@ async def validate_password_not_breached( """Raise ``ProxyException`` (400) if ``password`` appears in a known data breach. Fails open: an unreachable or misbehaving HIBP allows the password.""" - check_enabled: Final = general_settings.get("password_policy_check_breached_passwords", True) is not False - if not check_enabled: + if not await is_password_breached(password, general_settings, client): return - if not await _is_password_breached(password, client if client is not None else _hibp_client()): - return - raise ProxyException( - message=( - "This password appears in known data breaches and cannot be used. Please choose a different password." - ), - type=ProxyErrorTypes.validation_error, - param="password", - code=400, + raise breached_password_error() + + +def _strength_verdict(password: str, general_settings: Mapping[str, object]) -> ProxyException | None: + try: + validate_password_policy(password, general_settings) + except ProxyException as e: + return e + return None + + +async def validate_passwords_bulk( + passwords: Sequence[str], + general_settings: Mapping[str, object], + client: AsyncHTTPHandler | None = None, +) -> Mapping[str, ProxyException | None]: + """Per-unique-password policy verdicts for a batch: the ProxyException to + surface, or None when the password is acceptable. + + Deduplicates first, then issues every needed HIBP lookup concurrently, so a + batch caller pays one HIBP timeout window in the worst case instead of one + per password (each lookup still fails open independently).""" + unique_passwords: Final = tuple(dict.fromkeys(passwords)) + strength_verdicts: Final[Mapping[str, ProxyException | None]] = MappingProxyType( + {password: _strength_verdict(password, general_settings) for password in unique_passwords} + ) + to_screen: Final = tuple(password for password in unique_passwords if strength_verdicts[password] is None) + breached_flags: Final = await asyncio.gather( + *(is_password_breached(password, general_settings, client) for password in to_screen) + ) + breached_passwords: Final = frozenset(password for password, breached in zip(to_screen, breached_flags) if breached) + return MappingProxyType( + { + password: breached_password_error() if password in breached_passwords else strength_verdicts[password] + for password in unique_passwords + } ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ddef2127945..58a8ba0b09d 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -27,9 +27,14 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object -from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy +from litellm.proxy.auth.password_policy import ( + validate_password_not_breached, + validate_password_policy, + validate_passwords_bulk, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import ( @@ -162,11 +167,17 @@ def _team_membership_table( return team_membership_table -async def _hash_password_in_dict(data: dict, general_settings: Mapping[str, object]) -> None: - """Validate and hash password field in-place if present.""" +async def _hash_password_in_dict( + data: dict, general_settings: Mapping[str, object], password_prevalidated: bool = False +) -> None: + """Validate and hash password field in-place if present. + + ``password_prevalidated`` skips the policy checks for callers that already + validated the password (the bulk path screens its whole batch upfront).""" if "password" in data and data["password"] is not None: - validate_password_policy(data["password"], general_settings) - await validate_password_not_breached(data["password"], general_settings) + if not password_prevalidated: + validate_password_policy(data["password"], general_settings) + await validate_password_not_breached(data["password"], general_settings) data["password"] = hash_password(data["password"]) @@ -1429,6 +1440,7 @@ async def _update_single_user_helper( user_request: UpdateUserRequest, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None = None, + password_prevalidated: bool = False, ) -> dict[str, Any]: """ Helper function to update a single user. @@ -1451,7 +1463,7 @@ async def _update_single_user_helper( data_json: Final[dict] = user_request.model_dump(exclude_unset=True) non_default_values = _update_internal_user_params(data_json=data_json, data=user_request) - await _hash_password_in_dict(non_default_values, general_settings) + await _hash_password_in_dict(non_default_values, general_settings, password_prevalidated=password_prevalidated) existing_user_row: BaseModel | None = None if user_request.user_id: @@ -1700,19 +1712,38 @@ async def bulk_update_processed_users( users_to_update: list[UpdateUserRequest], user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None = None, + hibp_client: AsyncHTTPHandler | None = None, ) -> BulkUpdateUserResponse: + from litellm.proxy.proxy_server import general_settings + results: Final[list[UserUpdateResult]] = [] successful_updates = 0 failed_updates = 0 + # Screen the batch's passwords upfront and concurrently: done per-user + # inside the loop below, each HIBP lookup would be awaited serially and a + # degraded-slow HIBP could stretch a full batch to minutes, timing out the + # request after some updates already persisted. + password_verdicts: Final = await validate_passwords_bulk( + tuple(u.password for u in users_to_update if u.password is not None), + general_settings, + client=hibp_client, + ) + # Process each user update independently try: for user_request in users_to_update: try: + if ( + user_request.password is not None + and (password_error := password_verdicts.get(user_request.password)) is not None + ): + raise password_error response = await _update_single_user_helper( user_request=user_request, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, + password_prevalidated=True, ) # Record success results.append( diff --git a/tests/test_litellm/proxy/auth/test_password_policy.py b/tests/test_litellm/proxy/auth/test_password_policy.py index edc88d21218..f9f5025b57f 100644 --- a/tests/test_litellm/proxy/auth/test_password_policy.py +++ b/tests/test_litellm/proxy/auth/test_password_policy.py @@ -7,6 +7,7 @@ The breach-check (HIBP) tests inject a real AsyncHTTPHandler wrapping an httpx.MockTransport, so no network is touched and nothing is monkeypatched. """ +import asyncio import hashlib import httpx @@ -21,6 +22,7 @@ from litellm.proxy.auth.password_policy import ( get_password_policy, validate_password_not_breached, validate_password_policy, + validate_passwords_bulk, ) STRONG_PASSWORD = "Str0ng!Passw0rd" @@ -268,3 +270,76 @@ async def test_breach_check_fails_open_on_malformed_response_body(): client=_client_returning(f"{_sha1_upper('password12345')[5:]}:not-a-number"), ) assert result is None + + +@pytest.mark.asyncio +async def test_validate_passwords_bulk_screens_concurrently(): + """All HIBP lookups for a batch must be in flight at once: each handler + call stalls until every expected request has arrived, and a handler that + gives up waiting reports the password as breached. Serial awaiting (the + old per-user behavior) leaves each earlier request waiting forever for the + later ones, so every verdict comes back as a breach and the test fails.""" + passwords = ("Uniqu3!Passw0rd-a", "Uniqu3!Passw0rd-b", "Uniqu3!Passw0rd-c") + suffix_by_prefix = {_sha1_upper(p)[:5]: _sha1_upper(p)[5:] for p in passwords} + all_arrived = asyncio.Event() + arrivals: list[str] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + arrivals.append(request.url.path) + if len(arrivals) == len(passwords): + all_arrived.set() + try: + await asyncio.wait_for(all_arrived.wait(), timeout=5) + except TimeoutError: + return httpx.Response(200, text=f"{suffix_by_prefix[request.url.path.rsplit('/', 1)[-1]]}:1") + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + verdicts = await validate_passwords_bulk(passwords, {}, client=_client_with_transport(handler)) + assert set(arrivals) == {f"/range/{prefix}" for prefix in suffix_by_prefix} + assert all(verdicts[p] is None for p in passwords) + + +@pytest.mark.asyncio +async def test_validate_passwords_bulk_deduplicates_lookups(): + """500 users sharing one password must cost exactly one HIBP lookup.""" + password = "Sh@red-Passw0rd!" + request_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal request_count + request_count += 1 + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + verdicts = await validate_passwords_bulk((password,) * 500, {}, client=_client_with_transport(handler)) + assert request_count == 1 + assert verdicts == {password: None} + + +@pytest.mark.asyncio +async def test_validate_passwords_bulk_mixed_verdicts(): + """Weak passwords are rejected without an HIBP lookup; breached ones get + the breach error; acceptable ones map to None.""" + breached = "Br3ached!Passw0rd" + clean = "Cl3an!!Passw0rd42" + weak = "short1!" + breached_sha1 = _sha1_upper(breached) + looked_up_prefixes: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + looked_up_prefixes.append(request.url.path.rsplit("/", 1)[-1]) + if request.url.path == f"/range/{breached_sha1[:5]}": + return httpx.Response(200, text=f"{breached_sha1[5:]}:99") + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + verdicts = await validate_passwords_bulk((breached, clean, weak), {}, client=_client_with_transport(handler)) + assert _sha1_upper(weak)[:5] not in looked_up_prefixes + assert verdicts[clean] is None + assert "data breaches" in verdicts[breached].message + assert verdicts[breached].code == "400" + assert "12 characters" in verdicts[weak].message + + +@pytest.mark.asyncio +async def test_validate_passwords_bulk_empty_batch_makes_no_lookups(): + verdicts = await validate_passwords_bulk((), {}, client=_client_never_called()) + assert verdicts == {} diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index e7b5172e0fb..024d7b8300f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -11,6 +11,7 @@ from fastapi.testclient import TestClient from fastapi import HTTPException from pytest_mock import MockerFixture +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( LiteLLM_UserTableFiltered, @@ -67,9 +68,7 @@ async def test_ui_view_users_with_null_email(mocker, caplog): # Proxy admin: no org filter, no get_user_object call response = await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN), user_id="test_user", user_email=None, team_id=None, @@ -77,9 +76,7 @@ async def test_ui_view_users_with_null_email(mocker, caplog): page_size=50, ) - assert response == [ - LiteLLM_UserTableFiltered(user_id="test-user-null-email", user_email=None) - ] + assert response == [LiteLLM_UserTableFiltered(user_id="test-user-null-email", user_email=None)] @pytest.mark.asyncio @@ -103,9 +100,7 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker): mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), user_id=None, user_email="foo", team_id=None, @@ -128,9 +123,7 @@ async def test_ui_view_users_org_admin_filtered_by_org(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -268,9 +261,7 @@ async def test_ui_view_users_flag_on_team_admin_org_team(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -401,9 +392,7 @@ async def test_ui_view_users_flag_on_team_admin_org_member_no_team_id(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -462,9 +451,7 @@ async def test_ui_view_users_flag_on_team_admin_not_in_org_resolves_via_key_team async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -507,9 +494,7 @@ async def test_ui_view_users_flag_on_team_admin_not_in_org_resolves_via_key_team # No team_id query param, but team_id on the API key response = await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="team-admin-no-org", user_role=None, team_id=tid - ), + user_api_key_dict=UserAPIKeyAuth(user_id="team-admin-no-org", user_role=None, team_id=tid), user_id=None, user_email="u", team_id=None, @@ -538,13 +523,9 @@ def test_user_daily_activity_types(): # Assert all fields in SpendMetrics are reported in DailySpendMetadata as "total_" for field in spend_metrics.__dict__: if field.startswith("total_"): - assert hasattr( - daily_spend_metadata, field - ), f"Field {field} is not reported in DailySpendMetadata" + assert hasattr(daily_spend_metadata, field), f"Field {field} is not reported in DailySpendMetadata" else: - assert not hasattr( - daily_spend_metadata, field - ), f"Field {field} is reported in DailySpendMetadata" + assert not hasattr(daily_spend_metadata, field), f"Field {field} is reported in DailySpendMetadata" @pytest.mark.asyncio @@ -591,9 +572,7 @@ async def test_get_users_includes_timestamps(mocker): # Call get_users function directly with proxy admin auth admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - response = await get_users( - page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None - ) + response = await get_users(page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) print("user /list response: ", response) @@ -654,14 +633,10 @@ async def test_get_users_redacts_scim_enterprise_metadata(mocker): ) admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - response = await get_users( - page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None - ) + response = await get_users(page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) listed = response["users"][0] - assert listed.metadata == { - "scim_metadata": {"givenName": "Jane", "familyName": "Doe"} - } + assert listed.metadata == {"scim_metadata": {"givenName": "Jane", "familyName": "Doe"}} assert "scim_enterprise" not in (listed.metadata or {}) @@ -853,9 +828,7 @@ async def test_new_user_license_over_limit(mocker): mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check) # Create test request data - user_request = NewUserRequest( - user_email="test@example.com", user_role="internal_user" - ) + user_request = NewUserRequest(user_email="test@example.com", user_role="internal_user") # Mock user_api_key_dict mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin") @@ -916,9 +889,7 @@ async def test_new_user_license_gate_counts_only_billable_users(mocker): request = NewUserRequest(user_role="internal_user") # 2 active + 3 deactivated -> billable 2, not over max_users 2: gate passes - mocker.patch( - "litellm.proxy.proxy_server.prisma_client", _prisma(total=5, deactivated=3) - ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", _prisma(total=5, deactivated=3)) with pytest.raises(ProxyException) as passed: await new_user(data=request, user_api_key_dict=admin) assert key_gen.call_count == 1 @@ -926,9 +897,7 @@ async def test_new_user_license_gate_counts_only_billable_users(mocker): # 3 active, 0 deactivated -> billable 3, over max_users 2: gate blocks key_gen.reset_mock() - mocker.patch( - "litellm.proxy.proxy_server.prisma_client", _prisma(total=3, deactivated=0) - ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", _prisma(total=3, deactivated=0)) with pytest.raises(ProxyException) as blocked: await new_user(data=request, user_api_key_dict=admin) assert blocked.value.code == 403 or blocked.value.code == "403" @@ -978,14 +947,10 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check) # Test Case 1: INTERNAL_USER trying to create PROXY_ADMIN - user_request = NewUserRequest( - user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN - ) + user_request = NewUserRequest(user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN) # Mock user_api_key_dict with non-admin role - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER) # Call new_user function and expect ProxyException with pytest.raises(ProxyException) as exc_info: @@ -993,9 +958,7 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): # Verify the exception details assert exc_info.value.code == 403 or exc_info.value.code == "403" - assert "Only proxy admins can create administrative users" in str( - exc_info.value.message - ) + assert "Only proxy admins can create administrative users" in str(exc_info.value.message) assert "proxy_admin" in str(exc_info.value.message) assert "proxy_admin_viewer" in str(exc_info.value.message) assert str(LitellmUserRoles.PROXY_ADMIN) in str(exc_info.value.message) @@ -1008,15 +971,11 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): ) with pytest.raises(ProxyException) as exc_info2: - await new_user( - data=user_request_viewer, user_api_key_dict=mock_user_api_key_dict - ) + await new_user(data=user_request_viewer, user_api_key_dict=mock_user_api_key_dict) # Verify the exception details assert exc_info2.value.code == 403 or exc_info2.value.code == "403" - assert "Only proxy admins can create administrative users" in str( - exc_info2.value.message - ) + assert "Only proxy admins can create administrative users" in str(exc_info2.value.message) assert str(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) in str(exc_info2.value.message) @@ -1055,9 +1014,7 @@ async def test_new_user_non_admin_permissions_non_empty_rejected(mocker): user_role=LitellmUserRoles.INTERNAL_USER, permissions={"get_spend_routes": True}, ) - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(ProxyException) as exc_info: await new_user(data=data, user_api_key_dict=caller) @@ -1101,9 +1058,7 @@ async def test_new_user_non_admin_permissions_explicit_empty_rejected(mocker): permissions={}, ) assert "permissions" in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(ProxyException) as exc_info: await new_user(data=data, user_api_key_dict=caller) @@ -1156,9 +1111,7 @@ async def test_new_user_non_admin_omits_permissions_succeeds(mocker): user_role=LitellmUserRoles.INTERNAL_USER, ) assert "permissions" not in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) result = await new_user(data=data, user_api_key_dict=caller) assert result is not None @@ -1232,14 +1185,10 @@ async def test_update_single_user_non_admin_permissions_rejected(mocker): user_id="alice", permissions={"get_spend_routes": True}, ) - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc_info: - await _update_single_user_helper( - user_request=data, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=data, user_api_key_dict=caller) assert exc_info.value.status_code == 403 assert "permissions" in str(exc_info.value.detail) @@ -1261,14 +1210,10 @@ async def test_update_single_user_non_admin_permissions_explicit_empty_rejected( data = UpdateUserRequest(user_id="alice", permissions={}) assert "permissions" in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc_info: - await _update_single_user_helper( - user_request=data, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=data, user_api_key_dict=caller) assert exc_info.value.status_code == 403 assert "permissions" in str(exc_info.value.detail) @@ -1324,15 +1269,11 @@ async def test_user_info_url_encoding_plus_character(mocker): mock_request.url.query = "user_id=machine-user+alp-air-admin-b58-b@tempus.com" # Mock user_api_key_dict - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_admin", user_role="proxy_admin" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin", user_role="proxy_admin") # Call user_info function with the URL-decoded user_id (as FastAPI would pass it) # FastAPI would normally convert + to space, but our fix should handle this - decoded_user_id = ( - "machine-user alp-air-admin-b58-b@tempus.com" # What FastAPI gives us - ) + decoded_user_id = "machine-user alp-air-admin-b58-b@tempus.com" # What FastAPI gives us expected_user_id = "machine-user+alp-air-admin-b58-b@tempus.com" response = await user_info( @@ -1383,9 +1324,7 @@ async def test_user_info_nonexistent_user(mocker): mock_request = mocker.MagicMock(spec=Request) # Mock user_api_key_dict - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_admin", user_role="proxy_admin" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin", user_role="proxy_admin") # Call user_info function with a non-existent user_id nonexistent_user_id = "nonexistent-user@example.com" @@ -1423,14 +1362,10 @@ async def test_user_info_no_user_id_view_only_admin_gets_proxy_admin_payload(moc mock_get_user_info_for_proxy_admin, ) - viewer = UserAPIKeyAuth( - user_id="viewer", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value - ) + viewer = UserAPIKeyAuth(user_id="viewer", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value) mock_request = mocker.MagicMock(spec=Request) - response = await user_info( - user_id=None, user_api_key_dict=viewer, request=mock_request - ) + response = await user_info(user_id=None, user_api_key_dict=viewer, request=mock_request) mock_get_user_info_for_proxy_admin.assert_awaited_once_with(user_api_key_dict=viewer) assert response is admin_payload @@ -1457,9 +1392,7 @@ async def test_new_user_default_teams_flow(mocker): mock_prisma_client.db.litellm_usertable.count = mock_count persisted_user_row = mocker.MagicMock() persisted_user_row.teams = ["96fed65b-0182-4ff4-8429-2721cd7d42af"] - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - return_value=persisted_user_row - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=persisted_user_row) # Mock duplicate checks to pass async def mock_check_duplicate_user_email(*args, **kwargs): @@ -1527,26 +1460,20 @@ async def test_new_user_default_teams_flow(mocker): ) # Create test request data WITHOUT teams (teams should come from defaults) - user_request = NewUserRequest( - user_email="test@example.com", user_role="internal_user" - ) + user_request = NewUserRequest(user_email="test@example.com", user_role="internal_user") # Mock user_api_key_dict mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin") # Call new_user function - response = await new_user( - data=user_request, user_api_key_dict=mock_user_api_key_dict - ) + response = await new_user(data=user_request, user_api_key_dict=mock_user_api_key_dict) # Verify generate_key_helper_fn was called WITHOUT teams mock_generate_key_helper_fn.assert_called_once() call_kwargs = mock_generate_key_helper_fn.call_args.kwargs # Teams should be removed from the data passed to generate_key_helper_fn - assert ( - "teams" not in call_kwargs - ), "Teams should not be passed to generate_key_helper_fn" + assert "teams" not in call_kwargs, "Teams should not be passed to generate_key_helper_fn" assert call_kwargs["request_type"] == "user" assert call_kwargs["user_email"] == "test@example.com" assert call_kwargs["user_role"] == "internal_user" @@ -1591,24 +1518,16 @@ def test_update_internal_new_user_params_proxy_admin_role(): try: # Create test data with PROXY_ADMIN role - data = NewUserRequest( - user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN.value - ) + data = NewUserRequest(user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN.value) data_json = data.model_dump(exclude_unset=True) # Call the function result = _update_internal_new_user_params(data_json=data_json, data=data) # Assertions - default params should NOT be applied for PROXY_ADMIN - assert ( - "max_budget" not in result - ), "Default max_budget should NOT be applied to PROXY_ADMIN" - assert ( - "models" not in result - ), "Default models should NOT be applied to PROXY_ADMIN" - assert ( - "tpm_limit" not in result - ), "Default tpm_limit should NOT be applied to PROXY_ADMIN" + assert "max_budget" not in result, "Default max_budget should NOT be applied to PROXY_ADMIN" + assert "models" not in result, "Default models should NOT be applied to PROXY_ADMIN" + assert "tpm_limit" not in result, "Default tpm_limit should NOT be applied to PROXY_ADMIN" # These should still work assert result["user_email"] == "admin@example.com" @@ -1722,15 +1641,9 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): user_email_clause = where_clause.get("user_email", {}) # Check that the query structure is correct for case insensitive search - assert ( - "equals" in user_email_clause - ), "Query should use 'equals' for case insensitive search" - assert ( - user_email_clause.get("mode") == "insensitive" - ), "Query should use 'insensitive' mode" - assert ( - user_email_clause.get("equals") == "user@example.com" - ), "Query should search for the provided email" + assert "equals" in user_email_clause, "Query should use 'equals' for case insensitive search" + assert user_email_clause.get("mode") == "insensitive", "Query should use 'insensitive' mode" + assert user_email_clause.get("equals") == "user@example.com", "Query should search for the provided email" return mock_existing_user # Return existing user to simulate duplicate @@ -1741,9 +1654,7 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): await _check_duplicate_user_email("user@example.com", mock_prisma_client) assert exc_info.value.status_code == 409 - assert "User with email User@Example.com already exists" in str( - exc_info.value.detail - ) + assert "User with email User@Example.com already exists" in str(exc_info.value.detail) # Test Case 2: No duplicate found async def mock_find_first_no_duplicate(*args, **kwargs): @@ -1768,9 +1679,7 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): pytest.fail(f"Should not raise exception when no duplicate found, but got: {e}") # Test Case 3: None email should not cause issues - await _check_duplicate_user_email( - None, mock_prisma_client - ) # Should not raise exception + await _check_duplicate_user_email(None, mock_prisma_client) # Should not raise exception @pytest.mark.asyncio @@ -1880,9 +1789,7 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): # Verify dashboard key is not in results result_team_ids = [key.get("team_id") for key in result] - assert ( - UI_SESSION_TOKEN_TEAM_ID not in result_team_ids - ), "Dashboard key should be filtered out" + assert UI_SESSION_TOKEN_TEAM_ID not in result_team_ids, "Dashboard key should be filtered out" # Verify regular keys are included assert "regular-team" in result_team_ids, "Regular team key should be included" @@ -1892,9 +1799,7 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): result_tokens = [key.get("token") for key in result] assert "sk-regular-token" in result_tokens, "Regular key should be included" assert "sk-no-team-token" in result_tokens, "No-team key should be included" - assert ( - "sk-dashboard-token" not in result_tokens - ), "Dashboard key should not be included" + assert "sk-dashboard-token" not in result_tokens, "Dashboard key should not be included" def test_process_keys_for_user_info_handles_none_keys(monkeypatch): @@ -2331,9 +2236,7 @@ async def test_get_user_daily_activity_non_admin_cannot_view_other_users(monkeyp ) assert exc_info.value.status_code == 403 - assert "Non-admin users can only view their own spend data" in str( - exc_info.value.detail - ) + assert "Non-admin users can only view their own spend data" in str(exc_info.value.detail) # Case 2: Non-admin omits user_id — should default to their own user_id mock_response = MagicMock() @@ -2620,39 +2523,23 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): async def mock_find_unique(*args, **kwargs): return mock_user_row - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock find_many for teams (no teams) - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) # Mock all delete_many calls - mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock( - return_value=1 - ) - mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=1) + mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Call delete_user data = DeleteUserRequest(user_ids=["admin-creator"]) - user_api_key_dict = UserAPIKeyAuth( - user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + user_api_key_dict = UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN) await delete_user(data=data, user_api_key_dict=user_api_key_dict) @@ -2661,9 +2548,7 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): call_kwargs = mock_prisma_client.db.litellm_invitationlink.delete_many.call_args where_clause = call_kwargs.kwargs.get("where") or call_kwargs[1].get("where") - assert ( - "OR" in where_clause - ), "Should use OR to match user_id, created_by, and updated_by" + assert "OR" in where_clause, "Should use OR to match user_id, created_by, and updated_by" or_conditions = where_clause["OR"] assert len(or_conditions) == 3, "Should have 3 OR conditions" @@ -2706,9 +2591,7 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): async def mock_find_unique(*args, **kwargs): return mock_target_user - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Caller (org_admin_user) administers org-A. caller_membership = mocker.MagicMock() @@ -2734,16 +2617,12 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): return [caller_membership] return [] - mock_prisma_client.db.litellm_organizationmembership.find_many = mocker.AsyncMock( - side_effect=mock_find_memberships - ) + mock_prisma_client.db.litellm_organizationmembership.find_many = mocker.AsyncMock(side_effect=mock_find_memberships) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) data = DeleteUserRequest(user_ids=["victim"]) - user_api_key_dict = UserAPIKeyAuth( - user_id="org_admin_user", user_role=LitellmUserRoles.ORG_ADMIN - ) + user_api_key_dict = UserAPIKeyAuth(user_id="org_admin_user", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc: await delete_user(data=data, user_api_key_dict=user_api_key_dict) @@ -2751,11 +2630,8 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): # Critical: no delete_many calls should have executed. assert ( - not hasattr( - mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls" - ) - or len(mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls) - == 0 + not hasattr(mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls") + or len(mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls) == 0 ) @@ -2774,9 +2650,7 @@ async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker): mock_prisma_client = mocker.MagicMock() # user_email lookup yields None → would silently create pre-fix. - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=None) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest( @@ -2790,9 +2664,7 @@ async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker): ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=org_admin - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=org_admin) assert exc.value.status_code == 404 @@ -2836,17 +2708,13 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -2900,17 +2768,13 @@ async def test_user_info_v2_redacts_scim_enterprise_metadata(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -2919,9 +2783,7 @@ async def test_user_info_v2_redacts_scim_enterprise_metadata(mocker): ) assert isinstance(response, UserInfoV2Response) - assert response.metadata == { - "scim_metadata": {"givenName": "Jane", "familyName": "Doe"} - } + assert response.metadata == {"scim_metadata": {"givenName": "Jane", "familyName": "Doe"}} assert "scim_enterprise" not in (response.metadata or {}) @@ -2990,17 +2852,13 @@ async def test_user_info_v2_internal_user_can_query_self(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="self-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="self-user", user_role=LitellmUserRoles.INTERNAL_USER) response = await user_info_v2( request=mock_request, @@ -3035,17 +2893,13 @@ async def test_user_info_v2_internal_user_cannot_query_other(mocker): return mock_caller_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="caller-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="caller-user", user_role=LitellmUserRoles.INTERNAL_USER) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3092,17 +2946,13 @@ async def test_user_info_v2_no_user_id_defaults_to_self(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="my-user-id", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="my-user-id", user_role=LitellmUserRoles.INTERNAL_USER) # Call without user_id response = await user_info_v2( @@ -3130,17 +2980,13 @@ async def test_user_info_v2_nonexistent_user_returns_404(mocker): async def mock_find_unique(*args, **kwargs): return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3188,17 +3034,13 @@ async def test_user_info_v2_response_shape(mocker): async def mock_find_unique(*args, **kwargs): return mock_user_row - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -3233,9 +3075,7 @@ async def test_user_info_v2_response_shape(mocker): # The dashboard's user edit form hydrates its per-model budget rows from # these two, so dropping them makes a save replace the user's budgets. - assert response_dict["model_max_budget"] == { - "gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"} - } + assert response_dict["model_max_budget"] == {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}} assert response_dict["model_max_budget_usage"] == { "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"} } @@ -3294,9 +3134,7 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): return mock_target return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock team with caller as admin mock_team = mocker.MagicMock() @@ -3313,17 +3151,13 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): async def mock_find_many_teams(*args, **kwargs): return [mock_team] - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - side_effect=mock_find_many_teams - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(side_effect=mock_find_many_teams) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - team_admin_key = UserAPIKeyAuth( - user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + team_admin_key = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER) response = await user_info_v2( request=mock_request, @@ -3363,9 +3197,7 @@ async def test_user_info_v2_team_admin_cannot_query_non_team_member(mocker): return mock_target return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock team where caller is admin mock_team = mocker.MagicMock() @@ -3381,17 +3213,13 @@ async def test_user_info_v2_team_admin_cannot_query_non_team_member(mocker): async def mock_find_many_teams(*args, **kwargs): return [mock_team] - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - side_effect=mock_find_many_teams - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(side_effect=mock_find_many_teams) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - team_admin_key = UserAPIKeyAuth( - user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + team_admin_key = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3441,18 +3269,14 @@ async def test_user_info_v2_url_encoding_plus_character(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) mock_request.url.query = f"user_id={expected_user_id}" - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) # Simulate FastAPI converting + to space decoded_user_id = "machine-user admin@example.com" @@ -3549,9 +3373,7 @@ def test_enforce_user_info_access_admin_bypass(): _enforce_user_info_access, ) - admin = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN.value - ) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN.value) # Should not raise even when querying a different user _enforce_user_info_access(user_id="someone_else", user_api_key_dict=admin) @@ -3590,9 +3412,7 @@ def test_enforce_user_info_access_owner_allowed(): _enforce_user_info_access, ) - user = UserAPIKeyAuth( - user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value - ) + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) _enforce_user_info_access(user_id="alice", user_api_key_dict=user) @@ -3604,9 +3424,7 @@ def test_enforce_user_info_access_no_user_id_allowed(): _enforce_user_info_access, ) - user = UserAPIKeyAuth( - user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value - ) + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) _enforce_user_info_access(user_id=None, user_api_key_dict=user) @@ -3663,9 +3481,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budge "max_budget": 100, } existing_user.user_id = "user-1" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest.model_validate({"user_id": "user-1", budget_field: budget_value}) @@ -3675,9 +3491,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budge ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=caller) assert exc.value.status_code == 403 assert budget_field in str(exc.value.detail) mock_prisma_client.update_data.assert_not_called() @@ -3699,9 +3513,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_spend(mocker): "spend": 50.0, } existing_user.user_id = "user-1" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest( @@ -3714,9 +3526,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_spend(mocker): ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=caller) assert exc.value.status_code == 403 assert "spend" in str(exc.value.detail) @@ -3735,12 +3545,8 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): "max_budget": 100, } existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user", "max_budget": 500} - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user", "max_budget": 500}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3754,9 +3560,7 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): user_role=LitellmUserRoles.PROXY_ADMIN, ) - result = await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + result = await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) assert result is not None @@ -3772,12 +3576,8 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user", "spend": -25.0} - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user", "spend": -25.0}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3792,13 +3592,9 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): # without raising the recurring budget ceiling. Future changes should # continue allowing negative spend counters. user_request = UpdateUserRequest(user_id="target-user", spend=-25) - admin_caller = UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") @@ -3815,9 +3611,7 @@ async def test_user_update_rejects_non_finite_spend(mocker): existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mock_prisma_client.update_data = mocker.AsyncMock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3827,14 +3621,10 @@ async def test_user_update_rejects_non_finite_spend(mocker): ) user_request = UpdateUserRequest(user_id="target-user", spend=float("nan")) - admin_caller = UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) assert exc.value.status_code == 400 mock_prisma_client.update_data.assert_not_called() mock_invalidate.assert_not_awaited() @@ -3854,9 +3644,7 @@ async def test_resolve_user_email_metadata_maps_page_user_ids_to_email(mocker): mock_prisma_client = mocker.MagicMock() find_many = mocker.AsyncMock( return_value=[ - SimpleNamespace( - user_id="u1", user_email="alice@example.com", user_alias="Alice" - ), + SimpleNamespace(user_id="u1", user_email="alice@example.com", user_alias="Alice"), SimpleNamespace(user_id="u2", user_email=None, user_alias="bob-alias"), ] ) @@ -4055,19 +3843,13 @@ def _object_permission_mocks(mocker, existing_object_permission_id=None): } existing_user.user_id = "target-user" existing_user.object_permission_id = existing_object_permission_id - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = mocker.AsyncMock(return_value=None) mock_prisma_client.db.litellm_objectpermissiontable.upsert = mocker.AsyncMock( return_value=SimpleNamespace(object_permission_id="perm-new") ) mock_prisma_client.db.litellm_mcpservertable.find_many = mocker.AsyncMock(return_value=[]) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user"} - ) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user"}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -4103,9 +3885,7 @@ async def test_user_update_persists_mcp_entitlement_and_links_it(mocker): "mcp_tool_permissions": {"github": ["list_issues"]}, }, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) upsert_kwargs = mock_prisma_client.db.litellm_objectpermissiontable.upsert.call_args.kwargs @@ -4139,9 +3919,7 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker): user_id="target-user", object_permission={"mcp_tool_permissions": {"github": []}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} @@ -4174,9 +3952,7 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker): await _update_single_user_helper( user_request=UpdateUserRequest(user_id="target-user", object_permission={}), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) written = mock_prisma_client.update_data.call_args.kwargs["data"] @@ -4213,9 +3989,7 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock user_id="target-user", object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} @@ -4246,9 +4020,7 @@ async def test_non_admin_cannot_clear_their_own_mcp_entitlement(mocker): with pytest.raises(HTTPException) as exc: await _update_single_user_helper( user_request=UpdateUserRequest(user_id="target-user", object_permission={}), - user_api_key_dict=UserAPIKeyAuth( - user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER - ), + user_api_key_dict=UserAPIKeyAuth(user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER), ) assert exc.value.status_code == 403 @@ -4276,9 +4048,7 @@ async def test_non_admin_cannot_rewrite_their_own_mcp_entitlement(mocker): user_id="target-user", object_permission={"mcp_servers": [], "mcp_tool_permissions": {}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER - ), + user_api_key_dict=UserAPIKeyAuth(user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER), ) assert exc.value.status_code == 403 @@ -4295,9 +4065,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): return_value=SimpleNamespace(object_permission_id="perm-created") ) mock_prisma_client.db.litellm_mcpservertable.find_many = mocker.AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable.count = mocker.AsyncMock(return_value=0) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -4306,9 +4074,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): ) mock_generate = mocker.patch( "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", - new=mocker.AsyncMock( - return_value={"user_id": "new-human", "token": "sk-x", "expires": None} - ), + new=mocker.AsyncMock(return_value={"user_id": "new-human", "token": "sk-x", "expires": None}), ) mocker.patch( "litellm.proxy.hooks.user_management_event_hooks.UserManagementEventHooks.async_user_created_hook", @@ -4320,9 +4086,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): user_id="new-human", object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) created = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] @@ -4364,16 +4128,12 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker): response = await user_info_v2( request=SimpleNamespace(query_params={}), user_id="human-1", - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) assert response.object_permission is not None assert response.object_permission.mcp_servers == ["github"] - assert response.object_permission.mcp_tool_permissions == { - "github": ["list_issues"] - } + assert response.object_permission.mcp_tool_permissions == {"github": ["list_issues"]} @pytest.mark.asyncio @@ -4389,9 +4149,7 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker): ], ids=["supplied", "omitted", "empty"], ) -async def test_user_new_persists_model_max_budget( - monkeypatch, model_max_budget, expected_written -): +async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget, expected_written): """ /user/new used to echo model_max_budget back while writing {} to the user row, so a per-model budget looked configured and was read by nothing. @@ -4553,3 +4311,90 @@ async def test_user_update_rejects_breached_password(_admin_prisma): assert exc_info.value.code == "400" assert "data breaches" in exc_info.value.message _admin_prisma.db.litellm_usertable.find_first.assert_not_called() + + +def _hibp_client_with_handler(handler) -> AsyncHTTPHandler: + """A real AsyncHTTPHandler over httpx.MockTransport (the DI seam used + throughout test_password_policy.py), so no network is touched.""" + http_handler = AsyncHTTPHandler() + http_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + return http_handler + + +@pytest.mark.asyncio +async def test_bulk_update_breached_password_fails_only_that_user(_admin_prisma, mocker): + """In a bulk batch, a breached password fails only its own entry, before + any DB write for it; sibling entries with acceptable passwords persist.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + bulk_update_processed_users, + ) + + breached = "Br3ached!Passw0rd" + clean = "NewP@ssw0rd123" + breached_sha1 = hashlib.sha1(breached.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == f"/range/{breached_sha1[:5]}": + return httpx.Response(200, text=f"{breached_sha1[5:]}:1387") + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + mock_prisma_client = _admin_prisma + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "user-clean"} + existing_user.user_id = "user-clean" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-clean"}) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + response = await bulk_update_processed_users( + users_to_update=[ + UpdateUserRequest(user_id="user-breached", password=breached), + UpdateUserRequest(user_id="user-clean", password=clean), + ], + user_api_key_dict=admin_caller, + hibp_client=_hibp_client_with_handler(handler), + ) + + assert response.successful_updates == 1 + assert response.failed_updates == 1 + by_user = {r.user_id: r for r in response.results} + assert by_user["user-breached"].success is False + assert "data breaches" in by_user["user-breached"].error + assert by_user["user-clean"].success is True + (write_call,) = mock_prisma_client.update_data.call_args_list + assert write_call.kwargs["user_id"] == "user-clean" + + +@pytest.mark.asyncio +async def test_bulk_update_screens_shared_password_with_single_lookup(_admin_prisma, mocker): + """A batch where every user gets the same password costs one HIBP lookup, + not one per user (the serial per-user checks this regresses against).""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + bulk_update_processed_users, + ) + + lookup_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal lookup_count + lookup_count += 1 + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + mock_prisma_client = _admin_prisma + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "user-0"} + existing_user.user_id = "user-0" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-0"}) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + response = await bulk_update_processed_users( + users_to_update=[UpdateUserRequest(user_id=f"user-{i}", password="NewP@ssw0rd123") for i in range(5)], + user_api_key_dict=admin_caller, + hibp_client=_hibp_client_with_handler(handler), + ) + + assert response.successful_updates == 5 + assert lookup_count == 1 From 1d18d11fcf69388787b824b4d519e9855eede23e Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Tue, 8 Sep 2026 16:41:36 +0200 Subject: [PATCH 027/246] Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 58a8ba0b09d..4f037ffd43d 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -574,7 +574,7 @@ async def new_user( # generate_key_helper_fn only forwards object_permission_id, so without this the entitlement # the caller sent would be dropped on the floor. data_json = await _set_object_permission(data_json=data_json, prisma_client=prisma_client) - data_json.pop("password", None) # always None: NewUserRequest.password_not_supported rejects any other value + data_json.pop("password", None) teams = data.teams if teams is None: teams = check_if_default_team_set() From 5bb2c9e76f565bb4c4fb82ec527841a6340862e2 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Wed, 9 Sep 2026 13:46:29 +0200 Subject: [PATCH 028/246] fix(auth): annotate the strict-rule suppressions the merged gates now count The staging merge brought BLE001 into the strict ruff set and lowered the LIT002 ceiling, so the HIBP fail-open except and the params/headers dicts in password_policy.py now need their noqa and mutable-ok reasons. The headers dict moves to an annotated Final so the suppression fits the line limit. --- litellm/proxy/auth/password_policy.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index c27f276252a..7f06a0993d3 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -110,7 +110,7 @@ def validate_password_policy(password: str, general_settings: Mapping[str, objec def _hibp_client() -> AsyncHTTPHandler: return get_async_httpx_client( llm_provider=httpxSpecialProvider.PasswordBreachCheck, - params={"timeout": HIBP_TIMEOUT_SECONDS}, + params={"timeout": HIBP_TIMEOUT_SECONDS}, # mutable-ok: callee takes a bare dict (PEP 589) ) @@ -125,14 +125,18 @@ def _is_suffix_in_range_response(response_body: str, hash_suffix: str) -> bool: async def _is_password_breached(password: str, client: AsyncHTTPHandler) -> bool: # usedforsecurity=False: SHA-1 is only a lookup key into the HIBP dataset, so no security property rests on it sha1_hex: Final = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + headers: Final = { # mutable-ok: callee takes a bare dict (PEP 589) + "Add-Padding": "true", + "User-Agent": f"litellm-proxy/{version}", + } try: response: Final = await client.get( f"{HIBP_RANGE_API_BASE}/{sha1_hex[:5]}", - headers={"Add-Padding": "true", "User-Agent": f"litellm-proxy/{version}"}, + headers=headers, ) response.raise_for_status() breached: Final = _is_suffix_in_range_response(response.text, sha1_hex[5:]) - except Exception as e: + except Exception as e: # noqa: BLE001 # fail-open: any HIBP failure skips the check, never breaks the caller verbose_proxy_logger.warning("Breached-password check skipped, HIBP lookup failed: %s", e) return False return breached From d79a893e3776923ac33f802975e01e2557e130a5 Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Fri, 4 Sep 2026 09:02:34 +0200 Subject: [PATCH 029/246] feat(auth): add self-service change-password endpoint Admin password sets on /user/update and per-user /user/bulk_update stay supported and policy-enforced. The request model hides the password from repr so management alerts never format the plaintext, and the all_users bulk path rejects passwords instead of writing one plaintext value to every row. --- litellm/proxy/_types.py | 14 +- litellm/proxy/auth/route_checks.py | 17 +- .../internal_user_endpoints.py | 12 +- .../password_endpoints.py | 117 ++++++++ litellm/proxy/proxy_server.py | 4 + .../proxy/auth/test_route_checks.py | 213 ++++++------- .../test_internal_user_endpoints.py | 29 ++ .../test_password_endpoints.py | 283 ++++++++++++++++++ tests/test_litellm/proxy/test__types.py | 26 ++ .../ChangePasswordForm.integration.test.tsx | 68 +++++ .../change-password/ChangePasswordForm.tsx | 100 +++++++ .../app/(dashboard)/change-password/page.tsx | 7 + .../app/(dashboard)/hooks/useAuthorized.ts | 1 + .../Navbar/UserDropdown/UserDropdown.test.tsx | 52 +++- .../Navbar/UserDropdown/UserDropdown.tsx | 23 +- .../SidebarAccountMenu.test.tsx | 43 +++ .../SidebarAccountMenu/SidebarAccountMenu.tsx | 24 +- .../src/components/networking.tsx | 14 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 78 ++++- 19 files changed, 1004 insertions(+), 121 deletions(-) create mode 100644 litellm/proxy/management_endpoints/password_endpoints.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9b38b12ab07..7731ca736d5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -864,6 +864,7 @@ class LiteLLMRoutes(enum.Enum): "/organization/daily/activity", "/user/available_roles", # read-only role metadata; any authenticated user may read "/user/list", # org admins checked in endpoint; non-admins get 403 + "/user/password/change", # endpoint only ever writes the caller's own row "/model/{model_id}/update", "/prompt/list", "/prompt/info", @@ -1841,7 +1842,8 @@ class NewUserResponse(GenerateKeyResponse): class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest - password: str | None = None + # repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model + password: str | None = Field(default=None, repr=False) spend: float | None = None metadata: dict | None = None user_alias: str | None = None @@ -1871,6 +1873,16 @@ class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail): return values +class ChangePasswordRequest(LiteLLMPydanticObjectBase): + current_password: str = Field(repr=False) + new_password: str = Field(repr=False) + + +class ChangePasswordResponse(LiteLLMPydanticObjectBase): + user_id: str + message: str + + class DeleteUserRequest(LiteLLMPydanticObjectBase): user_ids: list[str] # required diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 953e3cf3e88..e4a36e73373 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -796,7 +796,8 @@ class RouteChecks: in the codebase is automatically readable by Admin Viewer without needing to remember to add it to an allowlist. 3. Unsafe HTTP method (POST/PUT/PATCH/DELETE): - - Allow `/user/update` only when restricted to user_email/password. + - Allow `/user/update` only when restricted to user_email. + - Allow `/user/password/change` (endpoint only writes the caller's own row). - Block all explicit writes in `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`. - Otherwise allow only if the route is in admin_viewer_routes / global_spend_tracking_routes (legacy explicit-allow set). @@ -816,10 +817,10 @@ class RouteChecks: if request_data is not None and isinstance(request_data, dict): _params_updated: Final = request_data.keys() for param in _params_updated: - if param not in ["user_email", "password"]: + if param != "user_email": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated", + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email can be updated", ) elif route in _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES or ( route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES) @@ -838,21 +839,25 @@ class RouteChecks: return # ── Unsafe HTTP method: explicit checks ────────────────────────── - # Allow `/user/update` for self-service email / password change. + # Allow `/user/update` for self-service email change. if route == "/user/update": if request_data is not None and isinstance(request_data, dict): for param in request_data: - if param not in ["user_email", "password"]: + if param != "user_email": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=( f"user not allowed to access this route, role= {_user_role}. " f"Trying to access: {route} and updating invalid param: {param}. " - "only user_email and password can be updated" + "only user_email can be updated" ), ) return + # Self-service password change; the endpoint only writes the caller's own row. + if route == "/user/password/change": + return + # Hard-block known write routes regardless of HTTP method (defensive # — these are POSTs in practice, but pinning them here protects # against future GET-shaped writes). diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 4f037ffd43d..368d5ed0c10 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1644,7 +1644,7 @@ async def user_update( Parameters: - user_id: Optional[str] - Specify a user id. If not set, a unique id will be generated. - user_email: Optional[str] - Specify a user email. - - password: Optional[str] - Specify a user password. + - password: Optional[str] - Set the user's password (admin only). Must satisfy the configured password policy. Users change their own password with POST /user/password/change. - user_alias: Optional[str] - A descriptive name for you to know who this user id refers to. - teams: Optional[list] - specify a list of team id's a user belongs to. - send_invite_email: Optional[bool] - Specify if an invite email should be sent. @@ -1881,6 +1881,16 @@ async def bulk_user_update( status_code=403, detail="Only proxy admins can update all users at once.", ) + if data.user_updates.password is not None: + raise HTTPException( + status_code=400, + detail={ + "error": ( + "Setting one password for all users is not supported. " + "Use per-user updates via the 'users' list instead." + ) + }, + ) # Optimized path for updating all users directly in database all_users_in_db: Final = await _user_table(prisma_client).find_many(order={"created_at": "desc"}) diff --git a/litellm/proxy/management_endpoints/password_endpoints.py b/litellm/proxy/management_endpoints/password_endpoints.py new file mode 100644 index 00000000000..78e53d7a777 --- /dev/null +++ b/litellm/proxy/management_endpoints/password_endpoints.py @@ -0,0 +1,117 @@ +""" +Self-service password management. + +/user/password/change + +Deliberately NOT wrapped in `management_endpoint_wrapper`: the wrapper emits +request kwargs to OTEL spans, which would log plaintext passwords. The audit +signal is emitted by hand below, with field names only, never values. +""" + +from typing import TYPE_CHECKING, Final + +from fastapi import APIRouter, Depends, HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + ChangePasswordRequest, + ChangePasswordResponse, + CommonProxyErrors, + LitellmTableNames, + UserAPIKeyAuth, +) +from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_helpers.audit_logs import create_object_audit_log +from litellm.proxy.utils import hash_password, verify_password +from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.user_repository import UserRepository + +if TYPE_CHECKING: + from prisma import models as prisma_models + + from litellm.proxy.utils import PrismaClient + +router: Final = APIRouter() + +_PASSWORD_CHANGED_AUDIT_VALUES: Final = '{"fields_changed": ["password"]}' + + +def _user_table( + prisma_client: "PrismaClient | None", +) -> "TableActions[prisma_models.LiteLLM_UserTable]": + user_table: Final[TableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table + return user_table + + +@router.post( + "/user/password/change", + tags=["Internal User management"], + dependencies=(Depends(user_api_key_auth),), +) +async def change_password( + data: ChangePasswordRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> ChangePasswordResponse: + """ + Change the calling user's own password. + + Requires the current password. The new password must satisfy the + configured password policy (`general_settings.password_policy_*`: minimum + length, character classes, and, when enabled, breached-password screening + via haveibeenpwned.com). + + Parameters: + - current_password: str - The user's current password. + - new_password: str - The password to change to. + """ + from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + user_id: Final = user_api_key_dict.user_id + if user_id is None: + raise HTTPException( + status_code=400, + detail={"error": "No user is associated with this session, so there is no password to change."}, + ) + + user_row: Final = await _user_table(prisma_client).find_first(where={"user_id": user_id}) + stored_password: Final = user_row.password if user_row is not None else None + if stored_password is None: + raise HTTPException( + status_code=400, + detail={ + "error": ( + "This account has no password set, so there is no password to change. " + "Passwords are set through an invitation link (POST /invitation/new)." + ) + }, + ) + + if not verify_password(data.current_password, stored_password): + raise HTTPException(status_code=400, detail={"error": "Current password is incorrect."}) + + validate_password_policy(data.new_password, general_settings) + await validate_password_not_breached(data.new_password, general_settings) + + await _user_table(prisma_client).update( + where={"user_id": user_id}, + data={"password": hash_password(data.new_password)}, + ) + + verbose_proxy_logger.info("Password changed via /user/password/change for user_id=%s", user_id) + await create_object_audit_log( + object_id=user_id, + action="updated", + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + table_name=LitellmTableNames.USER_TABLE_NAME, + after_value=_PASSWORD_CHANGED_AUDIT_VALUES, + ) + return ChangePasswordResponse(user_id=user_id, message="Password updated successfully.") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d2b486b4410..d803e8fd4be 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -549,6 +549,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( from litellm.proxy.management_endpoints.organization_endpoints import ( router as organization_router, ) +from litellm.proxy.management_endpoints.password_endpoints import ( + router as password_management_router, +) from litellm.proxy.management_endpoints.router_settings_endpoints import ( router as router_settings_router, ) @@ -18765,6 +18768,7 @@ app.include_router(pass_through_router) app.include_router(health_router) app.include_router(key_management_router) app.include_router(internal_user_router) +app.include_router(password_management_router) app.include_router(team_router) app.include_router(ui_sso_router) app.include_router(organization_router) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index c8b3d789665..4f58e3c86ff 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2,7 +2,6 @@ import os from datetime import datetime from unittest.mock import MagicMock, patch - import pytest from fastapi import HTTPException, Request @@ -38,7 +37,7 @@ def test_non_admin_config_update_route_rejected(): request.query_params = {} # Test that calling /config/update route raises HTTPException with 403 status - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -49,9 +48,8 @@ def test_non_admin_config_update_route_rejected(): ) # Verify the exception is raised with the correct message - assert ( - "Only proxy admin can be used to generate, delete, update info for new keys/users/teams" - in str(exc_info.value) + assert "Only proxy admin can be used to generate, delete, update info for new keys/users/teams" in str( + exc_info.value ) assert "Route=/config/update" in str(exc_info.value) assert "Your role=internal_user" in str(exc_info.value) @@ -130,7 +128,7 @@ def test_user_banner_update_rejected_for_non_admin(): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -703,9 +701,7 @@ def test_virtual_key_llm_api_route_includes_passthrough_prefix(route): valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["llm_api_routes"]) - result = RouteChecks.is_virtual_key_allowed_to_call_route( - route=route, valid_token=valid_token - ) + result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token) assert result is True @@ -730,9 +726,7 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route): valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["llm_api_routes"]) - result = RouteChecks.is_virtual_key_allowed_to_call_route( - route=route, valid_token=valid_token - ) + result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token) assert result is True @@ -802,18 +796,14 @@ def test_google_routes_with_dynamic_model_names_accessible_to_internal_users(): ) # If no exception is raised, the test passes except Exception as e: - pytest.fail( - f"Internal user should be able to access Google generateContent route. Got error: {str(e)}" - ) + pytest.fail(f"Internal user should be able to access Google generateContent route. Got error: {e!s}") def test_virtual_key_allowed_routes_with_multiple_litellm_routes_member_names(): """Test that virtual key works with multiple LiteLLMRoutes member names in allowed_routes""" # Create a UserAPIKeyAuth with multiple LiteLLMRoutes member names - valid_token = UserAPIKeyAuth( - user_id="test_user", allowed_routes=["openai_routes", "info_routes"] - ) + valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["openai_routes", "info_routes"]) # Test that routes from both groups are allowed result1 = RouteChecks.is_virtual_key_allowed_to_call_route( @@ -867,13 +857,9 @@ def test_virtual_key_allowed_routes_with_no_member_names_only_explicit(): ) # Test that explicit routes are allowed - result1 = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/chat/completions", valid_token=valid_token - ) + result1 = RouteChecks.is_virtual_key_allowed_to_call_route(route="/chat/completions", valid_token=valid_token) - result2 = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/custom/route", valid_token=valid_token - ) + result2 = RouteChecks.is_virtual_key_allowed_to_call_route(route="/custom/route", valid_token=valid_token) assert result1 is True assert result2 is True @@ -1241,9 +1227,7 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): ) assert exc_info.value.status_code == 403 - assert "Virtual key is not allowed to call this route" in str( - exc_info.value.detail - ) + assert "Virtual key is not allowed to call this route" in str(exc_info.value.detail) def test_check_passthrough_route_access_key_metadata_exact_match(): @@ -1702,9 +1686,7 @@ def test_videos_route_accessible_to_internal_users(): ) # If no exception is raised, the test passes except Exception as e: - pytest.fail( - f"Internal user should be able to access /v1/videos route. Got error: {str(e)}" - ) + pytest.fail(f"Internal user should be able to access /v1/videos route. Got error: {e!s}") def test_videos_route_with_virtual_key_llm_api_routes(): @@ -1726,12 +1708,8 @@ def test_videos_route_with_virtual_key_llm_api_routes(): ] for route in test_routes: - result = RouteChecks.is_virtual_key_allowed_to_call_route( - route=route, valid_token=valid_token - ) - assert ( - result is True - ), f"Virtual key with llm_api_routes should be able to access {route}" + result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token) + assert result is True, f"Virtual key with llm_api_routes should be able to access {route}" def test_non_proxy_admin_wildcard_allowed_routes(): @@ -1802,9 +1780,7 @@ def test_proxy_admin_viewer_can_access_global_spend_tags(): ) # If no exception is raised, the test passes except Exception as e: - pytest.fail( - f"proxy_admin_viewer should be able to access /global/spend/tags route. Got error: {str(e)}" - ) + pytest.fail(f"proxy_admin_viewer should be able to access /global/spend/tags route. Got error: {e!s}") # Routes returning proxy-wide spend across every team / customer / api_key. @@ -1832,7 +1808,7 @@ def test_internal_user_blocked_from_global_spend_routes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -1861,7 +1837,7 @@ def test_internal_user_view_only_blocked_from_global_spend_routes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, @@ -1963,9 +1939,7 @@ def test_proxy_admin_viewer_can_access_audit_logs(route): request_data={}, ) except Exception as e: - pytest.fail( - f"proxy_admin_viewer should be able to access {route} route. Got error: {str(e)}" - ) + pytest.fail(f"proxy_admin_viewer should be able to access {route} route. Got error: {e!s}") # ── Admin Viewer parity: Logs page endpoints ────────────────────────────────── @@ -2028,9 +2002,7 @@ def test_proxy_admin_viewer_can_access_logs_page_endpoints(route): request_data={}, ) except Exception as e: - pytest.fail( - f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}" - ) + pytest.fail(f"proxy_admin_viewer should be able to access {route}. Got error: {e!s}") @pytest.mark.parametrize( @@ -2140,7 +2112,7 @@ def test_internal_user_blocked_from_admin_viewer_logs_routes(route): if route not in INTERNAL_USER_BLOCKED_SUBSET: return - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2216,9 +2188,7 @@ def test_proxy_admin_viewer_can_access_settings_read_endpoints(route): request_data={}, ) except Exception as e: - pytest.fail( - f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}" - ) + pytest.fail(f"proxy_admin_viewer should be able to access {route}. Got error: {e!s}") # ── Admin Viewer parity: default-allow GET semantics ───────────────────────── @@ -2417,9 +2387,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: ) local_file = os.path.abspath(local_file) - spec = importlib.util.spec_from_file_location( - "local_enterprise_route_checks", local_file - ) + spec = importlib.util.spec_from_file_location("local_enterprise_route_checks", local_file) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod.EnterpriseRouteChecks @@ -2430,9 +2398,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2448,9 +2414,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2466,9 +2430,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2479,9 +2441,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks.should_call_route("/v1/chat/completions") assert exc_info.value.status_code == 403 - assert "LLM API routes are disabled for this instance." in str( - exc_info.value.detail - ) + assert "LLM API routes are disabled for this instance." in str(exc_info.value.detail) @patch("litellm.proxy.proxy_server.premium_user", True) def test_should_embeddings_still_blocked_when_llm_api_disabled(self): @@ -2489,9 +2449,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2509,9 +2467,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints: EnterpriseRouteChecks = self._get_enterprise_route_checks() with ( - patch.object( - EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=False - ), + patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=False), patch.object( EnterpriseRouteChecks, "is_management_routes_disabled", @@ -2530,9 +2486,7 @@ def test_route_in_additional_public_routes_wildcard_match(): from litellm.proxy.auth.auth_utils import route_in_additonal_public_routes with ( - patch( - "litellm.proxy.proxy_server.general_settings", {"public_routes": ["/api/*"]} - ), + patch("litellm.proxy.proxy_server.general_settings", {"public_routes": ["/api/*"]}), patch("litellm.proxy.proxy_server.premium_user", True), ): # Wildcard should match subpaths @@ -2624,7 +2578,7 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re ) # /config/update is still blocked - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2712,8 +2666,6 @@ def test_available_roles_accessible_to_non_admin_users(user_role): # ── _user_is_org_admin tests ────────────────────────────────────────────────── - - def _make_org_admin_user(org_id: str) -> LiteLLM_UserTable: membership = LiteLLM_OrganizationMembershipTable( user_id="org-admin-user", @@ -2836,9 +2788,7 @@ async def test_add_team_org_context_noop_when_org_id_already_present(): raise AssertionError("must not resolve when organization_id is present") body = {"team_id": "team-1", "organization_id": "org-explicit"} - out = await add_team_org_context_to_request_body( - route="/team/update", request_body=body, fetch_team_org_id=fetch - ) + out = await add_team_org_context_to_request_body(route="/team/update", request_body=body, fetch_team_org_id=fetch) assert out == body @@ -2850,9 +2800,7 @@ async def test_add_team_org_context_noop_for_other_routes(): raise AssertionError("must not resolve for a non-opted-in route") body = {"team_id": "team-1"} - out = await add_team_org_context_to_request_body( - route="/team/delete", request_body=body, fetch_team_org_id=fetch - ) + out = await add_team_org_context_to_request_body(route="/team/delete", request_body=body, fetch_team_org_id=fetch) assert out == body @@ -2865,9 +2813,7 @@ async def test_add_team_org_context_noop_when_team_has_no_org(): return None body = {"team_id": "team-1"} - out = await add_team_org_context_to_request_body( - route="/team/update", request_body=body, fetch_team_org_id=fetch - ) + out = await add_team_org_context_to_request_body(route="/team/update", request_body=body, fetch_team_org_id=fetch) assert out == body @@ -3151,9 +3097,7 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath(): # Removing the endpoint should clean up openai_routes # remove_endpoint_routes takes endpoint_id (UUID portion of # the route key "{id}:exact:{path}:{methods}") - registered = ( - InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() - ) + registered = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() endpoint_ids = {k.split(":")[0] for k in registered} for eid in endpoint_ids: InitPassThroughEndpointHelpers.remove_endpoint_routes(eid) @@ -3163,9 +3107,7 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath(): LiteLLMRoutes.openai_routes.value[:] = original_routes # Clean up any routes registered during this test to avoid # polluting the module-level _registered_pass_through_routes - registered = ( - InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() - ) + registered = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() for k in registered: InitPassThroughEndpointHelpers.remove_endpoint_routes(k.split(":")[0]) @@ -3196,8 +3138,7 @@ def test_provider_name_substring_not_classified_as_llm_route(route): from litellm.proxy.auth.route_checks import RouteChecks assert RouteChecks.is_llm_api_route(route=route) is False, ( - f"{route!r} should NOT be classified as an LLM API route — " - "provider-name substring match bypass" + f"{route!r} should NOT be classified as an LLM API route — provider-name substring match bypass" ) @@ -3219,9 +3160,7 @@ def test_legitimate_passthrough_routes_still_classified_as_llm_route(route): """Legitimate passthrough routes must still pass is_llm_api_route.""" from litellm.proxy.auth.route_checks import RouteChecks - assert ( - RouteChecks.is_llm_api_route(route=route) is True - ), f"{route!r} should be classified as an LLM API route" + assert RouteChecks.is_llm_api_route(route=route) is True, f"{route!r} should be classified as an LLM API route" @pytest.mark.parametrize( @@ -3279,7 +3218,7 @@ def test_internal_user_blocked_from_search_tool_writes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -3655,12 +3594,7 @@ def test_agent_inference_routes_stay_llm_api(route): def test_agent_routes_union_still_covers_both_halves(route): """Keys configured with allowed_routes=["agent_routes"] must keep both halves.""" - assert ( - RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.agent_routes.value - ) - is True - ) + assert RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_routes.value) is True @pytest.mark.parametrize("route", AGENT_MANAGEMENT_ROUTES) @@ -3714,6 +3648,75 @@ def test_agent_registry_route_gate_open_to_non_admin_roles(user_role, method, ro valid_token=valid_token, request_data={}, ) + + +def test_proxy_admin_viewer_user_update_password_param_rejected(): + """The self-service /user/update password carve-out is closed: non-admins + change their own password through /user/password/change, which verifies + the current password. Admin password sets don't pass through this check.""" + with pytest.raises(HTTPException) as exc_info: + RouteChecks._check_proxy_admin_viewer_access( + route="/user/update", + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + request_data={"password": "hunter2hunter2"}, + ) + assert exc_info.value.status_code == 403 + assert "password" in str(exc_info.value.detail) + + +def test_proxy_admin_viewer_user_update_user_email_still_allowed(): + request = MagicMock(spec=Request) + request.method = "POST" + + allowed = RouteChecks._check_proxy_admin_viewer_access( + route="/user/update", + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + request_data={"user_email": "viewer@example.com"}, + request=request, + ) + + assert allowed is None + + +def test_proxy_admin_viewer_can_change_own_password(): + request = MagicMock(spec=Request) + request.method = "POST" + + allowed = RouteChecks._check_proxy_admin_viewer_access( + route="/user/password/change", + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + request_data={"current_password": "a", "new_password": "b"}, + request=request, + ) + + assert allowed is None + + +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_non_admin_roles_can_change_own_password(user_role): + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = {} + + allowed = RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=LiteLLM_UserTable(user_id="test_user", user_role=user_role), + _user_role=user_role, + route="/user/password/change", + request=request, + valid_token=valid_token, + request_data={"current_password": "a", "new_password": "b"}, + ) + + assert allowed is None + + TEAM_CALLBACK_ROUTES = ( "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback", "/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback/langfuse", diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 024d7b8300f..a443235fe15 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4313,6 +4313,35 @@ async def test_user_update_rejects_breached_password(_admin_prisma): _admin_prisma.db.litellm_usertable.find_first.assert_not_called() +@pytest.mark.asyncio +async def test_bulk_update_all_users_rejects_a_password(_admin_prisma): + """The all_users fast path writes user_updates straight to update_many, + bypassing _update_single_user_helper. A password riding along would be + stored as unvalidated plaintext on every row, so it must be rejected + before any DB access.""" + from fastapi import HTTPException + + from litellm.proxy._types import UpdateUserRequestNoUserIDorEmail + from litellm.proxy.management_endpoints.internal_user_endpoints import bulk_user_update + from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( + BulkUpdateUserRequest, + ) + + data = BulkUpdateUserRequest( + all_users=True, + user_updates=UpdateUserRequestNoUserIDorEmail(password="Str0ng!Passw0rd"), + ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(HTTPException) as exc_info: + await bulk_user_update(data=data, user_api_key_dict=admin_caller) + + assert exc_info.value.status_code == 400 + assert "not supported" in str(exc_info.value.detail) + _admin_prisma.db.litellm_usertable.find_many.assert_not_called() + _admin_prisma.db.litellm_usertable.update_many.assert_not_called() + + def _hibp_client_with_handler(handler) -> AsyncHTTPHandler: """A real AsyncHTTPHandler over httpx.MockTransport (the DI seam used throughout test_password_policy.py), so no network is touched.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py new file mode 100644 index 00000000000..c7e0a385ae9 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py @@ -0,0 +1,283 @@ +""" +Tests for POST /user/password/change (litellm/proxy/management_endpoints/password_endpoints.py). + +HIBP traffic is intercepted with respx; no test here touches the network. +""" + +import hashlib +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +import respx +from fastapi import HTTPException + +from litellm.proxy._types import LitellmTableNames, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.management_endpoints.password_endpoints import change_password +from litellm.proxy.utils import hash_password, verify_password + +CURRENT_PASSWORD = "OldP@ssw0rd-2026" +NEW_PASSWORD = "NewP@ssw0rd-2026" + +_POLICY_NO_BREACH_CHECK = {"password_policy_check_breached_passwords": False} + + +def _make_user_row(password: str | None) -> MagicMock: + user = MagicMock() + user.user_id = "user-123" + user.password = password + return user + + +def _make_prisma(user: MagicMock | None) -> MagicMock: + prisma = MagicMock() + prisma.db.litellm_usertable.find_first = AsyncMock(return_value=user) + prisma.db.litellm_usertable.update = AsyncMock(return_value=user) + return prisma + + +def _caller(user_id: str | None = "user-123") -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id=user_id) + + +def _hibp_url_for(password: str) -> str: + sha1 = hashlib.sha1(password.encode(), usedforsecurity=False).hexdigest().upper() + return f"https://api.pwnedpasswords.com/range/{sha1[:5]}" + + +def _hibp_suffix_for(password: str) -> str: + return hashlib.sha1(password.encode(), usedforsecurity=False).hexdigest().upper()[5:] + + +@pytest.mark.asyncio +async def test_change_password_success_writes_new_scrypt_hash(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + response = await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert response.user_id == "user-123" + update_kwargs = prisma.db.litellm_usertable.update.call_args.kwargs + assert update_kwargs["where"] == {"user_id": "user-123"} + stored = update_kwargs["data"]["password"] + assert stored != NEW_PASSWORD + assert verify_password(NEW_PASSWORD, stored) + + +@pytest.mark.asyncio +async def test_change_password_rejects_wrong_current_password(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.status_code == 400 + assert "Current password is incorrect" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_change_password_rejects_session_without_user(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(user=None) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(user_id=None), + ) + + assert exc_info.value.status_code == 400 + prisma.db.litellm_usertable.find_first.assert_not_called() + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_change_password_rejects_account_without_password(): + """SSO users and the env-credential admin have no DB password row to change.""" + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(password=None)) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.status_code == 400 + assert "no password set" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_change_password_enforces_min_length(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(ProxyException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password="Short1!"), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.code == "400" + assert exc_info.value.type == ProxyErrorTypes.validation_error + assert exc_info.value.param == "password" + assert "at least 12 characters" in exc_info.value.message + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@respx.mock +async def test_change_password_rejects_breached_password(): + """With the default policy, the new password is screened against HIBP.""" + from litellm.proxy._types import ChangePasswordRequest + + breached_password = "Password123!" + respx.get(_hibp_url_for(breached_password)).mock( + return_value=httpx.Response(200, text=f"{_hibp_suffix_for(breached_password)}:1") + ) + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(ProxyException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=breached_password), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.code == "400" + assert exc_info.value.type == ProxyErrorTypes.validation_error + assert exc_info.value.param == "password" + assert "data breaches" in exc_info.value.message + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +@respx.mock +async def test_change_password_verifies_current_password_before_hibp_lookup(): + """A caller who fails current-password verification must not trigger any + HIBP traffic. The HIBP check fails open on errors, so an unmocked lookup + could not prove ordering; instead the route is registered and asserted + uncalled.""" + from litellm.proxy._types import ChangePasswordRequest + + hibp_route = respx.get(_hibp_url_for(NEW_PASSWORD)).mock(return_value=httpx.Response(200, text="")) + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.status_code == 400 + assert "Current password is incorrect" in exc_info.value.detail["error"] + assert not hibp_route.called + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_change_password_success_emits_redacted_audit_log(): + """A successful change must land in the audit trail as field names only; + the plaintext passwords must never reach the audit call.""" + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + audit_mock = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.management_endpoints.password_endpoints.create_object_audit_log", audit_mock), + ): + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + audit_mock.assert_awaited_once() + audit_kwargs = audit_mock.await_args.kwargs + assert audit_kwargs["object_id"] == "user-123" + assert audit_kwargs["action"] == "updated" + assert audit_kwargs["table_name"] == LitellmTableNames.USER_TABLE_NAME + assert audit_kwargs["after_value"] == '{"fields_changed": ["password"]}' + assert CURRENT_PASSWORD not in str(audit_kwargs) + assert NEW_PASSWORD not in str(audit_kwargs) + + +@pytest.mark.asyncio +async def test_change_password_failure_emits_no_audit_log(): + from litellm.proxy._types import ChangePasswordRequest + + prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD))) + audit_mock = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.management_endpoints.password_endpoints.create_object_audit_log", audit_mock), + ): + with pytest.raises(HTTPException): + await change_password( + data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + audit_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_change_password_requires_db(): + from litellm.proxy._types import ChangePasswordRequest + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + patch("litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK), # test-quality-ok: change_password reads proxy_server module globals; no injection seam + ): + with pytest.raises(HTTPException) as exc_info: + await change_password( + data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD), + user_api_key_dict=_caller(), + ) + + assert exc_info.value.status_code == 500 diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index 8f9c44d7a38..4ea30c7bf87 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -5,6 +5,7 @@ from pydantic import ValidationError from litellm.proxy._types import ( ROLES_WITHIN_ORG, + ChangePasswordRequest, GenerateKeyRequest, KeyRequest, LiteLLM_AuditLogs, @@ -293,3 +294,28 @@ def test_new_user_request_loudly_rejects_a_password(): def test_new_user_request_without_password_still_works(): request = NewUserRequest(user_email="alice@example.com") assert request.password is None + + +def test_update_user_request_accepts_a_password(): + """Admins set user passwords through /user/update; the value must survive + model validation so the endpoint can policy-check and hash it.""" + request = UpdateUserRequest(user_id="user-123", password="hunter2hunter2") + assert request.password == "hunter2hunter2" + + +def test_update_user_request_password_hidden_from_repr(): + """management_endpoint_wrapper string-formats endpoint kwargs into Slack + alerts, so the model's repr/str must never contain the plaintext password.""" + request = UpdateUserRequest(user_id="user-123", password="hunter2hunter2") + assert "hunter2hunter2" not in repr(request) + assert "hunter2hunter2" not in str(request) + + +def test_change_password_request_passwords_hidden_from_repr(): + """Any accidental str()/repr() of the request model (debug logs, exception + handlers, a future management_endpoint_wrapper) must never contain either + plaintext password.""" + request = ChangePasswordRequest(current_password="hunter2hunter2", new_password="NewP@ssw0rd-2026") + for rendered in (repr(request), str(request)): + assert "hunter2hunter2" not in rendered + assert "NewP@ssw0rd-2026" not in rendered diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx new file mode 100644 index 00000000000..e78f170cf0c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx @@ -0,0 +1,68 @@ +import { fireEvent, render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import ChangePasswordForm from "./ChangePasswordForm"; + +const mockChangePasswordCall = vi.fn(); +const mockToastSuccess = vi.fn(); + +vi.mock("@/components/networking", () => ({ + changePasswordCall: (...args: unknown[]) => mockChangePasswordCall(...args), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-session-token" }), +})); + +vi.mock("@/lib/toast", () => ({ + toast: { + success: (...args: unknown[]) => mockToastSuccess(...args), + fromError: vi.fn(), + }, +})); + +const fillForm = (values: { current: string; next: string; confirm: string }) => { + fireEvent.change(screen.getByLabelText("Current Password"), { target: { value: values.current } }); + fireEvent.change(screen.getByLabelText("New Password"), { target: { value: values.next } }); + fireEvent.change(screen.getByLabelText("Confirm New Password"), { target: { value: values.confirm } }); +}; + +const submit = () => fireEvent.click(screen.getByRole("button", { name: "Change Password" })); + +describe("ChangePasswordForm", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("sends the current and new password to the change endpoint and resets on success", async () => { + mockChangePasswordCall.mockResolvedValue({ user_id: "user-123", message: "Password updated successfully." }); + render(); + + fillForm({ current: "OldP@ssw0rd-2026", next: "NewP@ssw0rd-2026", confirm: "NewP@ssw0rd-2026" }); + submit(); + + expect(await screen.findByLabelText("Current Password")).toHaveValue(""); + expect(mockChangePasswordCall).toHaveBeenCalledWith("sk-session-token", "OldP@ssw0rd-2026", "NewP@ssw0rd-2026"); + expect(mockToastSuccess).toHaveBeenCalled(); + }); + + it("blocks submission when the confirmation does not match", async () => { + render(); + + fillForm({ current: "OldP@ssw0rd-2026", next: "NewP@ssw0rd-2026", confirm: "Different-2026" }); + submit(); + + expect(await screen.findByText("New passwords do not match")).toBeInTheDocument(); + expect(mockChangePasswordCall).not.toHaveBeenCalled(); + }); + + it("shows the proxy's rejection message unwrapped", async () => { + mockChangePasswordCall.mockRejectedValue(new Error("{'error': 'Current password is incorrect.'}")); + render(); + + fillForm({ current: "wrong-password", next: "NewP@ssw0rd-2026", confirm: "NewP@ssw0rd-2026" }); + submit(); + + expect(await screen.findByText("Current password is incorrect.")).toBeInTheDocument(); + expect(mockToastSuccess).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx new file mode 100644 index 00000000000..4c51b7f3d15 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx @@ -0,0 +1,100 @@ +"use client"; + +import React, { useState } from "react"; +import { CircleAlert } from "lucide-react"; +import { z } from "zod/v4"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; +import { PasswordInput } from "@/components/shared/PasswordInput"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent } from "@/components/ui/card"; +import { FieldGroup } from "@/components/ui/field"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { changePasswordCall } from "@/components/networking"; +import { extractProxyErrorMessage } from "@/lib/http/client"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { toast } from "@/lib/toast"; + +const changePasswordSchema = z + .object({ + currentPassword: z.string().min(1, "Current password is required"), + newPassword: z.string().min(1, "New password is required"), + confirmNewPassword: z.string().min(1, "Confirm your new password"), + }) + .refine((values) => values.newPassword === values.confirmNewPassword, { + message: "New passwords do not match", + path: ["confirmNewPassword"], + }); + +type ChangePasswordValues = z.infer; + +export function ChangePasswordForm() { + const { accessToken } = useAuthorized(); + const form = useZodForm(changePasswordSchema, { + defaultValues: { currentPassword: "", newPassword: "", confirmNewPassword: "" }, + }); + const [isPending, setIsPending] = useState(false); + const [submitError, setSubmitError] = useState(null); + + const handleSubmit = async (values: ChangePasswordValues) => { + if (!accessToken) return; + setSubmitError(null); + setIsPending(true); + try { + await changePasswordCall(accessToken, values.currentPassword, values.newPassword); + toast.success("Password updated"); + form.reset(); + } catch (error) { + setSubmitError(extractProxyErrorMessage(error)); + } finally { + setIsPending(false); + } + }; + + return ( +
+ + +

Change Password

+

+ Enter your current password and choose a new one. The new password must meet this proxy's password + policy. +

+ +
+ + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {submitError && ( + + + {submitError} + + )} + +
+ +
+
+ + +
+ ); +} + +export default ChangePasswordForm; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx new file mode 100644 index 00000000000..0a6ae926ceb --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx @@ -0,0 +1,7 @@ +"use client"; + +import ChangePasswordForm from "./ChangePasswordForm"; + +export default function ChangePasswordPage() { + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 40d1ec09d1f..089153cec76 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -50,6 +50,7 @@ const useAuthorized = () => { isViewOnly: isViewOnlySessionRole(decoded?.user_role), premiumUser: decoded?.premium_user ?? null, disabledPersonalKeyCreation: decoded?.disabled_non_admin_personal_key_creation ?? null, + loginMethod: decoded?.login_method ?? null, showSSOBanner: decoded?.login_method === "username_password", }; }; diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx index cad5ced340e..4bdf0da3407 100644 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx @@ -3,13 +3,25 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor } from "../../../../tests/test-utils"; import UserDropdown from "./UserDropdown"; -let mockUseAuthorizedImpl = () => ({ +let mockUseAuthorizedImpl: () => { + userId: string | null; + userEmail: string | null; + userRoleLabel: string; + premiumUser: boolean; + loginMethod?: string | null; +} = () => ({ userId: "test-user-id", userEmail: "test@example.com", userRoleLabel: "Admin", premiumUser: false, }); +const mockRouterPush = vi.fn(); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ push: mockRouterPush }), +})); + let mockUseDisableShowPromptsImpl = () => false; let mockGetLocalStorageItemImpl = (key: string): string | null => { @@ -143,6 +155,44 @@ describe("UserDropdown", () => { expect(mockOnLogout).toHaveBeenCalledTimes(1); }); + it("should navigate to the change-password page for username/password sessions", async () => { + mockUseAuthorizedImpl = () => ({ + userId: "test-user-id", + userEmail: "test@example.com", + userRoleLabel: "Admin", + premiumUser: false, + loginMethod: "username_password", + }); + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(getAccountTrigger()); + + await user.click(await screen.findByText("Change Password")); + + expect(mockRouterPush).toHaveBeenCalledWith(expect.stringContaining("change-password")); + }); + + it("should hide the change-password entry for SSO sessions", async () => { + mockUseAuthorizedImpl = () => ({ + userId: "test-user-id", + userEmail: "test@example.com", + userRoleLabel: "Admin", + premiumUser: false, + loginMethod: "sso", + }); + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(getAccountTrigger()); + + await waitFor(() => { + expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0); + }); + + expect(screen.queryByText("Change Password")).not.toBeInTheDocument(); + }); + it("should toggle hide new feature indicators switch", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx index 95f76dbb2cc..9c02defc778 100644 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx @@ -9,7 +9,9 @@ import { setLocalStorageItem, } from "@/utils/localStorageUtils"; import { navAccountDisplayName } from "@/components/Navbar/navDisplayName"; -import { ChevronDown, ChevronsUpDown, Crown, LogOut, Mail, ShieldCheck, User } from "lucide-react"; +import { uiHref } from "@/utils/uiHref"; +import { ChevronDown, ChevronsUpDown, Crown, KeyRound, LogOut, Mail, ShieldCheck, User } from "lucide-react"; +import { useRouter } from "next/navigation"; import { Avatar, AvatarFallback } from "@/components/ui/avatar"; import { Badge } from "@/components/ui/badge"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; @@ -63,7 +65,9 @@ interface UserDropdownProps { } const UserDropdown: React.FC = ({ onLogout, variant = "navbar", collapsed = false }) => { - const { userId, userEmail, userRoleLabel: userRole, premiumUser } = useAuthorized(); + const { userId, userEmail, userRoleLabel: userRole, premiumUser, loginMethod } = useAuthorized(); + const router = useRouter(); + const [open, setOpen] = useState(false); const disableShowPrompts = useDisableShowPrompts(); const disableBlogPosts = useDisableBlogPosts(); const disableBouncingIcon = useDisableBouncingIcon(); @@ -197,7 +201,7 @@ const UserDropdown: React.FC = ({ onLogout, variant = "navbar const displayName = navAccountDisplayName(userEmail, userId); return ( - + {variant === "sidebar" ? ( = ({ onLogout, variant = "navbar > {renderUserInfoSection()} + {loginMethod === "username_password" && ( + + )} + )} + + + + + ); +} + +describe("JEV classifier editor", () => { + afterEach(() => vi.mocked(useAuthorized).mockReset()); + it("uses built-in JEV without a license and preserves custom tiers and context through reload", () => { + renderWithProviders(); + expect(screen.getByLabelText("Classifier Model")).toBeInTheDocument(); + expect(screen.getByText("Reasoning Effort")).toBeInTheDocument(); + expect(screen.getByText("Classifier Prompt")).toBeInTheDocument(); + expect(screen.getByRole("switch", { name: "Use images for classification" })).toBeInTheDocument(); + fireEvent.click(screen.getByRole("radio", { name: /JEV Classifier/ })); + expect(screen.getByRole("tab", { name: "Complexity" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByLabelText("JEV Model")).toHaveValue("jev-latest"); + expect(screen.getByLabelText("JEV Instructions")).toBeDisabled(); + expect(screen.queryByLabelText("Classifier Model")).not.toBeInTheDocument(); + expect(screen.queryByText("Reasoning Effort")).not.toBeInTheDocument(); + expect(screen.queryByText("Classifier Prompt")).not.toBeInTheDocument(); + expect(screen.queryByRole("switch", { name: "Use images for classification" })).not.toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("JEV Model"), { target: { value: "jev-test" } }); + fireEvent.change(screen.getByLabelText("JEV Timeout (ms)"), { target: { value: "4200" } }); + fireEvent.change(screen.getByLabelText("Context Window Size"), { target: { value: "6" } }); + fireEvent.change(screen.getByLabelText("Circuit breaker cooldown (seconds)"), { target: { value: "50" } }); + fireEvent.click(screen.getByRole("switch", { name: "Classifier circuit breaker" })); + fireEvent.click(screen.getByRole("button", { name: "Customize tiers" })); + fireEvent.click(screen.getByRole("button", { name: "Save and reload" })); + expect(screen.getByRole("radio", { name: /JEV Classifier/ })).toBeChecked(); + expect(screen.getByLabelText("JEV Model")).toHaveValue("jev-test"); + expect(screen.getByLabelText("JEV Timeout (ms)")).toHaveValue(4200); + expect(screen.getByLabelText("Context Window Size")).toHaveValue("6"); + expect(screen.getByRole("switch", { name: "Classifier circuit breaker" })).not.toBeChecked(); + fireEvent.click(screen.getByRole("button", { name: "Probe current config" })); + expect(testAutoRouterRouting).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + complexity_router_config: expect.objectContaining({ + classifier_type: "jev", + jev_classifier_config: { + model: "jev-test", + timeout_ms: 4200, + circuit_breaker_enabled: false, + circuit_breaker_cooldown_seconds: 50, + }, + tiers: expect.objectContaining({ QUICK: ["fast"] }), + }), + }), + ); + }); + + it("allows licensed instructions and can restore built-in instructions", () => { + const authorized = useAuthorized(); + vi.mocked(useAuthorized).mockReturnValue({ ...authorized, premiumUser: true }); + const LicensedForm = () => { + const [value, setValue] = useState({ + ...initial, + classifier_type: "jev", + jev_classifier_config: { model: "jev-latest", timeout_ms: 3000, instructions: "Existing instructions" }, + }); + return ; + }; + renderWithProviders(); + expect(screen.getByLabelText("JEV Instructions")).toBeEnabled(); + fireEvent.change(screen.getByLabelText("JEV Instructions"), { target: { value: "New instructions" } }); + expect(screen.getByLabelText("JEV Instructions")).toHaveValue("New instructions"); + fireEvent.click(screen.getByRole("button", { name: "Restore built-in JEV instructions" })); + expect(screen.getByLabelText("JEV Instructions")).toHaveValue(""); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx new file mode 100644 index 00000000000..25286eaef07 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/JevClassifierConfig.tsx @@ -0,0 +1,88 @@ +import React, { useId } from "react"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Textarea } from "@/components/ui/textarea"; +import { SimpleTooltip } from "@/components/ui/tooltip"; +import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { defaultJevClassifierConfig } from "./jev_classifier_config"; + +export default function JevClassifierConfig({ + value, + onChange, +}: { + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +}) { + const id = useId(); + const { premiumUser } = useAuthorized(); + const config = value.jev_classifier_config ?? defaultJevClassifierConfig(); + const update = (patch: Partial) => + onChange({ ...value, jev_classifier_config: { ...config, ...patch } }); + + return ( +
+

+ Uses TypeSafe System One Choice evaluation with your configured tiers +

+
+ + update({ model: event.target.value })} /> +
+
+ + update({ timeout_ms: Number(event.target.value) })} + /> +
+ + update({ + circuit_breaker_enabled: next.circuit_breaker_enabled, + circuit_breaker_cooldown_seconds: next.circuit_breaker_cooldown_seconds, + }) + } + /> +
+ + +
+