From b93030f84e7a414d2106528114b09f1fca1ad1aa Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:20:25 +0000 Subject: [PATCH 01/13] fix(vertex_ai): surface real error/status on vertex batch create instead of IndexError 500 --- litellm/llms/vertex_ai/batches/handler.py | 44 ++++++++++++---- .../llms/vertex_ai/batches/transformation.py | 46 ++++++++++++++--- .../llms/vertex_ai/batches/test_handler.py | 50 +++++++++++++++---- .../vertex_ai/batches/test_transformation.py | 43 +++++++++++++++- 4 files changed, 152 insertions(+), 31 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index ada1356fb6b..f0fd5480c75 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -13,7 +13,7 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( @@ -98,7 +98,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -130,7 +132,9 @@ class VertexAIBatchPrediction(VertexLLM): ) raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -242,7 +246,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -292,7 +298,9 @@ class VertexAIBatchPrediction(VertexLLM): headers=headers, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -365,7 +373,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = ( @@ -390,7 +400,9 @@ class VertexAIBatchPrediction(VertexLLM): params=params, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = ( @@ -475,7 +487,9 @@ class VertexAIBatchPrediction(VertexLLM): raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) # HTTPHandler.get() does not accept a timeout parameter retrieve_response = sync_handler.get( @@ -488,7 +502,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -521,7 +538,9 @@ class VertexAIBatchPrediction(VertexLLM): ) raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) # AsyncHTTPHandler.get() does not accept a timeout parameter retrieve_response = await client.get( @@ -534,7 +553,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index df903ba7ef0..e4299bcf2a0 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,7 +1,9 @@ from typing import Any, Dict, Optional +from urllib.parse import unquote from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest @@ -199,16 +201,40 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ - from urllib.parse import unquote - - decoded_uri = unquote(gcs_file_uri) - - model_path = decoded_uri.split("publishers/")[1] - parts = model_path.split("/") - model = f"publishers/{'/'.join(parts[:3])}" + model = cls._parse_model_from_gcs_file(gcs_file_uri) + if model is None: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch creation requires the model to be part of `input_file_id`, but " + f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + "Either upload the input file through LiteLLM (POST /v1/files with " + "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " + "pass a uri of the form " + "gs:////publishers//models//" + ), + ) return model + @classmethod + def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: + """ + Returns the `publishers//models/` path from a gcs uri, or None if the uri + does not contain one. + """ + _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") + if not separator: + return None + + parts = model_path.split("/") + if len(parts) < 3 or parts[1] != "models" or not parts[2]: + return None + + return f"publishers/{'/'.join(parts[:3])}" + @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool: """ @@ -216,7 +242,11 @@ class VertexAIBatchTransformation: LiteLLM-managed unified file id) with a `publishers/` model path that `_get_model_from_gcs_file` can parse. """ - return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id + return ( + input_file_id is not None + and input_file_id.startswith("gs://") + and cls._parse_model_from_gcs_file(input_file_id) is not None + ) @classmethod def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str: diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index cacea234777..b9fb5dfe3c5 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -40,6 +40,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) +from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 HMOD = "litellm.llms.vertex_ai.batches.handler" @@ -184,7 +185,7 @@ def test_create_batch_sync_non_200_raises(): client.post.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500") as exc_info: h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -196,6 +197,32 @@ def test_create_batch_sync_non_200_raises(): max_retries=None, ) + assert exc_info.value.status_code == 500 + assert "error text" in str(exc_info.value) + + +def test_create_batch_input_file_id_without_model_raises_400_before_post(): + """A gs:// uri with no publishers//models/ path is a 400, not a bare 500.""" + h = _make_handler() + client = MagicMock() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data={"input_file_id": "gs://bucket/batch-input.jsonl"}, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + assert "gs://bucket/batch-input.jsonl" in str(exc_info.value) + client.post.assert_not_called() + def test_create_batch_async_non_200_raises(): h = _make_handler() @@ -216,9 +243,12 @@ def test_create_batch_async_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 403"): + with pytest.raises(VertexAIError, match="Error: 403") as exc_info: _run(coro) + assert exc_info.value.status_code == 403 + assert "error text" in str(exc_info.value) + # =========================================================================== # # retrieve_batch @@ -292,7 +322,7 @@ def test_retrieve_batch_sync_non_200_raises(): patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), patch(f"{HMOD}.safe_get", return_value=_http_response(status_code=404)), ): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.retrieve_batch( _is_async=False, batch_id=BATCH_ID, @@ -438,7 +468,7 @@ def test_list_batches_sync_non_200_raises(): client.get.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.list_batches( _is_async=False, after=None, @@ -530,7 +560,7 @@ def test_cancel_batch_sync_cancel_post_non_200_raises(): client.post.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -552,7 +582,7 @@ def test_cancel_batch_sync_retrieve_non_200_raises(): client.get.return_value = _http_response(status_code=404) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -672,7 +702,7 @@ def test_async_retrieve_batch_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -726,7 +756,7 @@ def test_async_list_batches_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -779,7 +809,7 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) async_client_post500.get.assert_not_awaited() @@ -801,5 +831,5 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): _run(coro) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 1b37ade6b30..8352ec16389 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -25,6 +25,7 @@ from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.utils import LiteLLMBatch # noqa: E402 @@ -69,6 +70,24 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +@pytest.mark.parametrize( + "input_file_id", + [ + "gs://bucket/no-model-here.jsonl", + "gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", + "gs://bucket/publishers/google/models", + "gs://bucket/publishers/google/models//file-uuid", + ], +) +def test_transform_openai_request_unparseable_model_raises_400(input_file_id: str): + """An input_file_id with no parseable model path is a client error, not an IndexError -> 500.""" + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": input_file_id}) + + assert exc_info.value.status_code == 400 + assert input_file_id in str(exc_info.value) + + # =========================================================================== # # transform_vertex_ai_batch_response_to_openai_batch_response # =========================================================================== # @@ -299,9 +318,29 @@ def test_get_model_from_gcs_file_url_encoded(): assert T._get_model_from_gcs_file(encoded) == "publishers/google/models/gemini-1.5-flash-001" -def test_get_model_from_gcs_file_no_publishers_raises(): - with pytest.raises(IndexError): +def test_get_model_from_gcs_file_no_publishers_raises_400(): + with pytest.raises(VertexAIError) as exc_info: T._get_model_from_gcs_file("gs://bucket/no-model-here.jsonl") + assert exc_info.value.status_code == 400 + + +# =========================================================================== # +# is_unmanaged_gcs_batch_input_file_id +# =========================================================================== # + + +@pytest.mark.parametrize( + "input_file_id, expected", + [ + (INPUT_FILE, True), + (None, False), + ("file-abc123", False), + ("gs://bucket/no-model-here.jsonl", False), + ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + ], +) +def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): + assert T.is_unmanaged_gcs_batch_input_file_id(input_file_id) is expected # =========================================================================== # From c5c5a276790529e2de3378654864fd847530c5a0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:19:42 +0000 Subject: [PATCH 02/13] fix(files): enforce require_managed_files on file retrieve, content and delete Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../openai_files_endpoints/common_utils.py | 30 +++++ .../openai_files_endpoints/files_endpoints.py | 7 ++ .../test_files_endpoint.py | 116 ++++++++++++++++++ 3 files changed, 153 insertions(+) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 87514b46dbd..3eef3868c94 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -866,6 +866,36 @@ def validate_managed_files_requirement( ) +def validate_managed_file_id_requirement(file_id: str) -> None: + """ + Enforce proxy-level managed files on the file read/delete routes when + ``litellm.require_managed_files`` is enabled. + + Ownership is only recorded for LiteLLM managed files, so a raw provider file id sent to + retrieve/content/delete is forwarded to the provider under shared credentials without any + tenant check; knowing another tenant's provider file id would be enough to read or delete it. + + Raises: + HTTPException: 400 if ``file_id`` is not a LiteLLM managed file id. + """ + import litellm + from fastapi import HTTPException + + if litellm.require_managed_files is not True: + return + + if _is_base64_encoded_unified_file_id(file_id): + return + + raise HTTPException( + status_code=400, + detail=( + "Raw provider file ids cannot be used when require_managed_files is enabled in " + "litellm_settings. Use the LiteLLM managed file id returned when the file was created." + ), + ) + + def _extract_model_param(request: "Request", request_body: dict) -> str | None: """ Extract model parameter from request. diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 4e4718272bd..37f1ced6996 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, + validate_managed_file_id_requirement, validate_managed_files_requirement, ) from litellm.proxy.utils import ProxyLogging, is_known_model @@ -612,6 +613,8 @@ async def get_file_content( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + # Include original request and headers in the data base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -908,6 +911,8 @@ async def get_file( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -1098,6 +1103,8 @@ async def delete_file( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ac01c6ae1d1..24b814bae1f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3051,3 +3051,119 @@ def test_list_files_key_allowed_openai_model_still_resolves_team_credentials( mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["team-gpt"] ) assert captured_kwargs.get("api_key") == "team-openai-key" + + +@pytest.mark.parametrize( + "http_method, url, patched_litellm_call", + [ + ("get", "/v1/files/file-victim-abc123", "litellm.afile_retrieve"), + ("get", "/v1/files/file-victim-abc123/content", "litellm.afile_content"), + ("delete", "/v1/files/file-victim-abc123", "litellm.afile_delete"), + ], +) +def test_require_managed_files_rejects_raw_provider_file_id( + mocker: MockerFixture, + monkeypatch, + llm_router: Router, + http_method: str, + url: str, + patched_litellm_call: str, +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", True) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_call = mocker.patch(patched_litellm_call, new=mocker.AsyncMock()) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="attacker-user" + ) + + try: + response = getattr(client, http_method)( + url, headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + monkeypatch.setattr("litellm.require_managed_files", False) + + assert response.status_code == 400, response.text + mock_call.assert_not_called() + + +def _unified_managed_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-3.5-turbo", "file-victim-abc123", "gpt-3.5-turbo-id" + ) + return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + +def test_require_managed_files_allows_unified_managed_file_id(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_file_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", True) + + validate_managed_file_id_requirement(file_id=_unified_managed_file_id()) + + +def test_managed_file_id_requirement_is_opt_in(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_file_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", False) + + validate_managed_file_id_requirement(file_id="file-victim-abc123") + + +def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", False) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_retrieve = mocker.patch( + "litellm.afile_retrieve", + new=mocker.AsyncMock( + return_value=OpenAIFileObject( + id="file-victim-abc123", + object="file", + bytes=3, + created_at=1234567890, + filename="test.txt", + purpose="user_data", + status="uploaded", + ) + ), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="some-user" + ) + + try: + response = client.get( + "/v1/files/file-victim-abc123", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + mock_retrieve.assert_called_once() From ede84eee15ab0703598f87a8fce2612e406d49e9 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 03:25:38 +0000 Subject: [PATCH 03/13] ci: give the remaining pull_request workflows a concurrency group Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/check-schema-sync.yml | 4 ++++ .github/workflows/conventional-commits.yml | 4 ++++ .github/workflows/guard-fork-dependencies.yml | 4 ++++ .github/workflows/helm_unit_test.yml | 4 ++++ .github/workflows/test-linting.yml | 4 ++++ .github/workflows/test-litellm-ui-build.yml | 4 ++++ .github/workflows/test-litellm-ui-lint.yml | 4 ++++ .github/workflows/test-mcp.yml | 4 ++++ .github/workflows/test-model-map.yaml | 4 ++++ 9 files changed, 36 insertions(+) diff --git a/.github/workflows/check-schema-sync.yml b/.github/workflows/check-schema-sync.yml index 0e5e2804e60..a4e78d2c44c 100644 --- a/.github/workflows/check-schema-sync.yml +++ b/.github/workflows/check-schema-sync.yml @@ -10,6 +10,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: check-sync: name: Verify schema.prisma copies match root diff --git a/.github/workflows/conventional-commits.yml b/.github/workflows/conventional-commits.yml index 69ade24d028..eb9eb69f8b6 100644 --- a/.github/workflows/conventional-commits.yml +++ b/.github/workflows/conventional-commits.yml @@ -14,6 +14,10 @@ on: permissions: pull-requests: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: lint-pr-title: name: Validate PR title diff --git a/.github/workflows/guard-fork-dependencies.yml b/.github/workflows/guard-fork-dependencies.yml index f4cbdd63cdf..6b366da78d4 100644 --- a/.github/workflows/guard-fork-dependencies.yml +++ b/.github/workflows/guard-fork-dependencies.yml @@ -15,6 +15,10 @@ on: permissions: {} +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: guard: name: Block fork dependency changes diff --git a/.github/workflows/helm_unit_test.yml b/.github/workflows/helm_unit_test.yml index a44d412c781..f95848945a0 100644 --- a/.github/workflows/helm_unit_test.yml +++ b/.github/workflows/helm_unit_test.yml @@ -9,6 +9,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: unit-test: runs-on: ubuntu-latest diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 5e333f2a3ca..3db3fb07a94 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: lint: runs-on: ubuntu-latest diff --git a/.github/workflows/test-litellm-ui-build.yml b/.github/workflows/test-litellm-ui-build.yml index 39f4bc1428a..618b0195b5a 100644 --- a/.github/workflows/test-litellm-ui-build.yml +++ b/.github/workflows/test-litellm-ui-build.yml @@ -10,6 +10,10 @@ on: - litellm_oss_staging - "litellm_**" +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: build-ui: runs-on: ubuntu-latest diff --git a/.github/workflows/test-litellm-ui-lint.yml b/.github/workflows/test-litellm-ui-lint.yml index ecc739a87e2..e03d89ee26a 100644 --- a/.github/workflows/test-litellm-ui-lint.yml +++ b/.github/workflows/test-litellm-ui-lint.yml @@ -10,6 +10,10 @@ on: - litellm_oss_staging - "litellm_**" +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: frontend-lint: runs-on: ubuntu-latest diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index a5a4e722133..05cc13d0af2 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: test: runs-on: ubuntu-latest diff --git a/.github/workflows/test-model-map.yaml b/.github/workflows/test-model-map.yaml index cf4b0eb21a1..c2770e5da4c 100644 --- a/.github/workflows/test-model-map.yaml +++ b/.github/workflows/test-model-map.yaml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: validate-model-prices-json: runs-on: ubuntu-latest From 557d14cc71f7a0c1b44fa36f20b0da4cf9330e47 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:27:00 -0700 Subject: [PATCH 04/13] fix(lint): make strict-gate noqas survive base ruff and flag stale ones --- litellm/utils.py | 2 +- ruff-strict-budget.json | 2 +- ruff-strict.toml | 8 ++++++++ ruff.toml | 4 +++- 4 files changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index cdf1cc3cf23..911de83b785 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5769,7 +5769,7 @@ def json_schema_type(python_type_name: str): return python_to_json_schema_types.get(python_type_name, "string") -def function_to_dict(input_function) -> dict: # noqa: C901 +def function_to_dict(input_function) -> dict: """Using type hints and numpy-styled docstring, produce a dictionary usable for OpenAI function calling diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 8ff4bfb36c0..a8af4eabb3f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -252,7 +252,7 @@ "limit": 67 }, "RUF100": { - "limit": 100 + "limit": 0 }, "S110": { "limit": 218 diff --git a/ruff-strict.toml b/ruff-strict.toml index 01faf04805f..d20fb2e4d7d 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -4,6 +4,14 @@ extend = "ruff.toml" preview = true select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"] extend-select = [] +# Overrides the inherited list: rules this gate enforces itself must NOT be external here, +# so this config's RUF100 flags their stale `# noqa` directives. What remains external is +# only what other tooling enforces: upstream litellm's ruff config, plus the base ruff.toml +# rules (T20, E7xx/F5xx/F8xx) this select list doesn't re-enable. +external = [ + "T20", "E731", "F541", "F841", + "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", +] [lint.mccabe] max-complexity = 15 diff --git a/ruff.toml b/ruff.toml index 095e3e24c52..00743e0f38a 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,9 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "C901", "TID251", + "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "FURB", "I001", "LOG015", "N999", "PERF", + "PIE", "PL", "PYI", "RET", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", + "RUF046", "RUF051", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] From c3536c29a0ebb8e5ad63663d6e0798f80eb5ac9c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:38:53 -0700 Subject: [PATCH 05/13] fix(lint): cover every base-owned ruff rule in the strict gate's external list --- ruff-strict.toml | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/ruff-strict.toml b/ruff-strict.toml index d20fb2e4d7d..974c49c787b 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -6,10 +6,13 @@ select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B0 extend-select = [] # Overrides the inherited list: rules this gate enforces itself must NOT be external here, # so this config's RUF100 flags their stale `# noqa` directives. What remains external is -# only what other tooling enforces: upstream litellm's ruff config, plus the base ruff.toml -# rules (T20, E7xx/F5xx/F8xx) this select list doesn't re-enable. +# only what other tooling enforces: every base ruff.toml rule this select list doesn't +# re-enable (all of the default E/F families plus T20/PGH004/RUF008/RUF009, minus the +# strict-selected F401 and RUF100; F4 is split out so stale F401 noqas stay detectable), +# plus upstream litellm's ruff config. external = [ - "T20", "E731", "F541", "F841", + "T20", "PGH004", "RUF008", "RUF009", "E4", "E7", "E9", + "F402", "F404", "F406", "F407", "F5", "F6", "F7", "F8", "F9", "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] From 7bffbbd1f2132a00206b6864ddd102cfade9a94f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:00:50 -0700 Subject: [PATCH 06/13] refactor(vertex_ai): drop unreachable post-path status checks in batches handler HTTPHandler.post and AsyncHTTPHandler.post call raise_for_status before returning, so the status_code != 200 branches after the create and cancel POSTs could never run. Non-2xx already surfaces as httpx.HTTPStatusError from inside the client. The checks after GETs stay: the get helpers return without raising. Tests that faked a non-raising POST response are replaced by HTTPStatusError propagation coverage. --- litellm/llms/vertex_ai/batches/handler.py | 22 +---- .../llms/vertex_ai/batches/test_handler.py | 92 +++---------------- 2 files changed, 16 insertions(+), 98 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 8e36c2a0faa..6481b67fad7 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -98,11 +98,6 @@ class VertexAIBatchPrediction(VertexLLM): data=json.dumps(vertex_batch_request), ) - if response.status_code != 200: - raise VertexAIError( - status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" - ) - _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( response=_json_response @@ -132,10 +127,6 @@ class VertexAIBatchPrediction(VertexLLM): error_body[:1000], ) raise - if response.status_code != 200: - raise VertexAIError( - status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" - ) _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -473,7 +464,7 @@ class VertexAIBatchPrediction(VertexLLM): sync_handler: Final = _get_httpx_client() try: - response: Final = sync_handler.post( + sync_handler.post( url=api_base, headers=headers, data=json.dumps({}), @@ -487,11 +478,6 @@ class VertexAIBatchPrediction(VertexLLM): ) raise - if response.status_code != 200: - raise VertexAIError( - status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" - ) - # HTTPHandler.get() does not accept a timeout parameter retrieve_response: Final = sync_handler.get( url=retrieve_api_base, @@ -525,7 +511,7 @@ class VertexAIBatchPrediction(VertexLLM): llm_provider=litellm.LlmProviders.VERTEX_AI, ) try: - response: Final = await client.post( + await client.post( url=api_base, headers=headers, data=json.dumps({}), @@ -538,10 +524,6 @@ class VertexAIBatchPrediction(VertexLLM): e.response.text[:1000], ) raise - if response.status_code != 200: - raise VertexAIError( - status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" - ) # AsyncHTTPHandler.get() does not accept a timeout parameter retrieve_response: Final = await client.get( diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index b9fb5dfe3c5..9535bf17411 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -5,8 +5,10 @@ The handler is HTTP/auth glue around the (separately-tested) pure ``VertexAIBatchTransformation``. Each public method (create / retrieve / list / cancel) resolves a Vertex access token + URL, branches on ``_is_async`` (returning the coroutine in the async case, doing the sync HTTP call otherwise), -checks the HTTP status, and parses the JSON into ``LiteLLMBatch`` (or the OpenAI -list shape). +and parses the JSON into ``LiteLLMBatch`` (or the OpenAI list shape). POST-backed +calls rely on the client's ``raise_for_status`` (non-2xx surfaces as +``httpx.HTTPStatusError``); GET-backed calls return without raising, so the +handler checks their status codes itself. We mock only true I/O / auth seams: * ``_ensure_access_token`` - the Vertex credential seam. Returns a fixed @@ -20,7 +22,7 @@ We mock only true I/O / auth seams: what URL/headers/body, and that the response is parsed into the litellm type. Sibling seams are asserted NOT called where relevant. -The ``_is_async`` branch, status-code error paths, and the cancel +The ``_is_async`` branch, the error paths, and the cancel retrieve-after-cancel sequencing run for real. """ @@ -179,13 +181,19 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): sync_client.post.assert_not_called() -def test_create_batch_sync_non_200_raises(): +def test_create_batch_sync_httpstatuserror_propagates(): + """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the + sync create path must surface that error, not swallow it.""" h = _make_handler() client = MagicMock() - client.post.return_value = _http_response(status_code=500) + request = httpx.Request("POST", "https://x/batchPredictionJobs") + err_response = httpx.Response(status_code=500, request=request, text="boom") + client.post.side_effect = httpx.HTTPStatusError( + "boom", request=request, response=err_response + ) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(VertexAIError, match="Error: 500") as exc_info: + with pytest.raises(httpx.HTTPStatusError): h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -197,9 +205,6 @@ def test_create_batch_sync_non_200_raises(): max_retries=None, ) - assert exc_info.value.status_code == 500 - assert "error text" in str(exc_info.value) - def test_create_batch_input_file_id_without_model_raises_400_before_post(): """A gs:// uri with no publishers//models/ path is a 400, not a bare 500.""" @@ -224,32 +229,6 @@ def test_create_batch_input_file_id_without_model_raises_400_before_post(): client.post.assert_not_called() -def test_create_batch_async_non_200_raises(): - h = _make_handler() - async_client = MagicMock() - async_client.post = AsyncMock(return_value=_http_response(status_code=403)) - - with ( - patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), - patch(f"{HMOD}.get_async_httpx_client", return_value=async_client), - ): - coro = h.create_batch( - _is_async=True, - create_batch_data=CREATE_DATA, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - with pytest.raises(VertexAIError, match="Error: 403") as exc_info: - _run(coro) - - assert exc_info.value.status_code == 403 - assert "error text" in str(exc_info.value) - - # =========================================================================== # # retrieve_batch # =========================================================================== # @@ -554,27 +533,6 @@ def test_cancel_batch_async_returns_coroutine_posts_then_retrieves(): assert post_kwargs["url"].endswith(":cancel") -def test_cancel_batch_sync_cancel_post_non_200_raises(): - h = _make_handler() - client = MagicMock() - client.post.return_value = _http_response(status_code=500) - - with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(VertexAIError, match="Error: 500"): - h.cancel_batch( - _is_async=False, - batch_id=BATCH_ID, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - # cancel POST failed -> retrieve GET must never fire - client.get.assert_not_called() - - def test_cancel_batch_sync_retrieve_non_200_raises(): h = _make_handler() client = MagicMock() @@ -791,28 +749,6 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): _run(coro) async_client.get.assert_not_awaited() - # (a2) cancel POST returns a plain non-200 (no exception) -> raises - async_client_post500 = MagicMock() - async_client_post500.post = AsyncMock(return_value=_http_response(status_code=500)) - async_client_post500.get = AsyncMock() - with ( - patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), - patch(f"{HMOD}.get_async_httpx_client", return_value=async_client_post500), - ): - coro = h.cancel_batch( - _is_async=True, - batch_id=BATCH_ID, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - with pytest.raises(VertexAIError, match="Error: 500"): - _run(coro) - async_client_post500.get.assert_not_awaited() - # (b) retrieve-after-cancel returns non-200 async_client2 = MagicMock() async_client2.post = AsyncMock(return_value=_http_response(json_body={})) From f304b7b19faa9743636965c79a9709fd5bd2b2d0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:10:33 -0700 Subject: [PATCH 07/13] refactor(lint): graduate the 35 zero-violation strict rules into ruff.toml Every strict-gate rule whose budget ceiling was already 0 moves into the base config's lint.extend-select, so editors and ruff check --fix surface the diagnostics directly and the budget file shrinks to rules with real debt. Graduates stay in ruff-strict.toml's select so the strict RUF100 pass keeps policing their stale noqa directives, and base external entries they made redundant (FURB, I001, RUF010, RUF022, RUF023, RUF051) are dropped so base RUF100 polices those directly. UP037 had two violations hidden behind a star import; importing Literal explicitly fixes them so UP037 can graduate too. New drift tests pin the invariants: every strict-selected rule is budgeted or hard-failed by base, every base-owned rule stays visible to exactly one RUF100 pass, and graduated rules fail the normal ruff run. --- .../internal_user_endpoints.py | 2 +- ruff-strict-budget.json | 105 --------- ruff.toml | 22 +- tests/test_litellm/test_ruff_strict_gate.py | 220 +++++++++++++++++- 4 files changed, 237 insertions(+), 112 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 640a735c916..abc5d3e53ff 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,7 +17,7 @@ import json import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, cast +from typing import Any, Final, Literal, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a8af4eabb3f..7e350c184af 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -56,9 +56,6 @@ "B026": { "limit": 3 }, - "B033": { - "limit": 0 - }, "BLE001": { "limit": 2924 }, @@ -113,18 +110,6 @@ "F401": { "limit": 17 }, - "FURB136": { - "limit": 0 - }, - "FURB168": { - "limit": 0 - }, - "FURB188": { - "limit": 0 - }, - "I001": { - "limit": 0 - }, "LOG015": { "limit": 5 }, @@ -137,18 +122,9 @@ "PERF401": { "limit": 12 }, - "PERF402": { - "limit": 0 - }, "PERF403": { "limit": 34 }, - "PIE790": { - "limit": 0 - }, - "PIE800": { - "limit": 0 - }, "PIE804": { "limit": 18 }, @@ -158,9 +134,6 @@ "PLC0206": { "limit": 26 }, - "PLC0208": { - "limit": 0 - }, "PLC0414": { "limit": 46 }, @@ -170,24 +143,12 @@ "PLR0206": { "limit": 1 }, - "PLR0402": { - "limit": 0 - }, "PLR1704": { "limit": 3 }, - "PLR1711": { - "limit": 0 - }, "PLR1714": { "limit": 257 }, - "PLR1730": { - "limit": 0 - }, - "PLR2044": { - "limit": 0 - }, "PLW0127": { "limit": 57 }, @@ -206,27 +167,12 @@ "PLW1510": { "limit": 2 }, - "PYI030": { - "limit": 0 - }, "PYI036": { "limit": 3 }, - "PYI041": { - "limit": 0 - }, - "PYI064": { - "limit": 0 - }, - "RET501": { - "limit": 0 - }, "RET504": { "limit": 177 }, - "RUF010": { - "limit": 0 - }, "RUF012": { "limit": 241 }, @@ -236,18 +182,9 @@ "RUF019": { "limit": 38 }, - "RUF022": { - "limit": 0 - }, - "RUF023": { - "limit": 0 - }, "RUF046": { "limit": 4 }, - "RUF051": { - "limit": 0 - }, "RUF059": { "limit": 67 }, @@ -272,18 +209,12 @@ "SIM113": { "limit": 3 }, - "SIM114": { - "limit": 0 - }, "SIM115": { "limit": 2 }, "SIM117": { "limit": 7 }, - "SIM118": { - "limit": 0 - }, "SIM201": { "limit": 1 }, @@ -302,9 +233,6 @@ "TC004": { "limit": 5 }, - "TC005": { - "limit": 0 - }, "TID251": { "limit": 1240 }, @@ -323,46 +251,13 @@ "TRY300": { "limit": 860 }, - "UP006": { - "limit": 0 - }, - "UP007": { - "limit": 0 - }, - "UP008": { - "limit": 0 - }, - "UP012": { - "limit": 0 - }, - "UP018": { - "limit": 0 - }, - "UP024": { - "limit": 0 - }, "UP028": { "limit": 2 }, "UP031": { "limit": 2 }, - "UP032": { - "limit": 0 - }, - "UP034": { - "limit": 0 - }, - "UP035": { - "limit": 0 - }, "UP036": { "limit": 1 - }, - "UP037": { - "limit": 0 - }, - "UP045": { - "limit": 0 } } diff --git a/ruff.toml b/ruff.toml index 00743e0f38a..9b90910b355 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,14 +1,26 @@ lint.ignore = ["F405", "E402", "F403"] -lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] +# The second group is the strict gate's graduates: rules the codebase already has zero +# violations of, so they hard-fail here instead of being ratcheted in ruff-strict-budget.json. +# That gives editors and `ruff check --fix` the diagnostic, which the gate script cannot. +lint.extend-select = [ + "T20", "PGH004", "RUF008", "RUF009", "RUF100", + "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208", + "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", + "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", "UP012", + "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", +] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external # so RUF100 leaves their directives alone: the strict gate (ruff-strict.toml) and upstream # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ - # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "FURB", "I001", "LOG015", "N999", "PERF", - "PIE", "PL", "PYI", "RET", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", - "RUF046", "RUF051", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", + # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml). + # Family entries whose every strict rule graduated into extend-select above (FURB), and + # standalone graduated codes (I001, RUF010, RUF022, RUF023, RUF051), are dropped so this + # config's RUF100 polices their directives itself. + "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "LOG015", "N999", "PERF", + "PIE", "PL", "PYI", "RET", "RUF012", "RUF015", "RUF019", + "RUF046", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/test_litellm/test_ruff_strict_gate.py index abdeb6feecc..206207acb09 100644 --- a/tests/test_litellm/test_ruff_strict_gate.py +++ b/tests/test_litellm/test_ruff_strict_gate.py @@ -1,16 +1,24 @@ import importlib.util +import json +import re +import shutil import subprocess +import sys +import tomllib from pathlib import Path import pytest -_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "ruff_strict_gate.py" +_REPO_ROOT = Path(__file__).resolve().parents[2] +_MODULE_PATH = _REPO_ROOT / "scripts" / "ruff_strict_gate.py" _spec = importlib.util.spec_from_file_location("ruff_strict_gate", _MODULE_PATH) gate = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(gate) Violation = gate.Violation +_ENABLED_BY_RUFF_DEFAULTS = frozenset({"F401"}) + def rule(name, limit): return {name: {"limit": limit}} @@ -151,3 +159,213 @@ def test_base_point_mid_merge_advances_to_the_merged_in_base_tip(tmp_path): repo, _, base_tip = _branched_repo(tmp_path) _git(repo, "merge", "--no-commit", "--no-ff", "main") assert gate.resolve_base_point("main", cwd=repo) == base_tip + + +def _lint_section(config_name: str) -> dict: + return tomllib.loads((_REPO_ROOT / config_name).read_text())["lint"] + + +def _base_external() -> tuple[str, ...]: + return tuple(_lint_section("ruff.toml")["external"]) + + +def _strict_external() -> tuple[str, ...]: + return tuple(_lint_section("ruff-strict.toml")["external"]) + + +def _strict_selected() -> frozenset: + return frozenset(_lint_section("ruff-strict.toml")["select"]) + + +def _prefix_covered(code: str, prefixes: tuple[str, ...]) -> bool: + return any(code.startswith(prefix) for prefix in prefixes) + + +def _selected_by_the_normal_config() -> frozenset: + return frozenset(_lint_section("ruff.toml")["extend-select"]) | _ENABLED_BY_RUFF_DEFAULTS + + +def _budgeted_rules() -> frozenset: + return frozenset(json.loads((_REPO_ROOT / "ruff-strict-budget.json").read_text())) + + +def _ruff_binary() -> str | None: + beside_interpreter = Path(sys.executable).with_name("ruff") + return str(beside_interpreter) if beside_interpreter.exists() else shutil.which("ruff") + + +_RUFF = _ruff_binary() +_needs_ruff = pytest.mark.skipif(_RUFF is None, reason="ruff is not installed in this environment") + + +def _ruff_output_for_noqa(code: str, *extra_args: str) -> str: + proc = subprocess.run( + [ + _RUFF, + "check", + "--no-cache", + "--stdin-filename", + "litellm/types/_external_probe.py", + *extra_args, + "-", + ], + cwd=_REPO_ROOT, + input=f"def _probe(x: int): # noqa: {code}\n return x\n", + capture_output=True, + text=True, + ) + return proc.stdout + + +def test_every_strict_gate_rule_is_protected_from_base_ruf100(): + unprotected = frozenset( + selector + for selector in _strict_selected() + if not _prefix_covered(selector, _base_external()) + and selector not in _selected_by_the_normal_config() + ) + assert unprotected == frozenset(), ( + f"`ruff check` deletes any `# noqa` naming {sorted(unprotected)} as unused, so suppressing " + "one of those strict-gate rules breaks lint. Cover them in ruff.toml's lint.external or " + "enable them in its lint.extend-select." + ) + + +def test_every_selected_rule_keeps_stale_noqa_detection_somewhere(): + policed_by_strict = frozenset( + selector + for selector in _strict_selected() + if not _prefix_covered(selector, _strict_external()) + ) + policed_by_base = frozenset( + selector + for selector in _selected_by_the_normal_config() + if not _prefix_covered(selector, _base_external()) + ) + shadowed = ( + _strict_selected() | _selected_by_the_normal_config() + ) - policed_by_strict - policed_by_base + assert shadowed == frozenset(), ( + f"no config's RUF100 can ever report a stale `# noqa` for {sorted(shadowed)}: every config " + "that selects each of them also shadows it with an external entry. Narrow the external " + "entry in ruff.toml or ruff-strict.toml." + ) + + +_BASE_OWNED_FAMILY = re.compile(r"E[479]\d+|F\d+|T20\d+") +_BASE_OWNED_SINGLES = frozenset({"PGH004", "RUF008", "RUF009", "RUF100"}) + + +@pytest.fixture(scope="module") +def all_ruff_rule_codes() -> frozenset: + listing = subprocess.run( + [_RUFF, "rule", "--all", "--output-format", "json"], + capture_output=True, + text=True, + ) + assert listing.returncode == 0, listing.stderr + return frozenset( + entry["code"] for entry in json.loads(listing.stdout) if "Removed" not in entry["status"] + ) + + +@_needs_ruff +def test_every_base_owned_rule_is_external_or_selected_in_the_strict_config(all_ruff_rule_codes): + base_owned = frozenset( + code + for code in all_ruff_rule_codes + if _BASE_OWNED_FAMILY.fullmatch(code) or code in _BASE_OWNED_SINGLES + ) + stranded = frozenset( + code + for code in base_owned + if code not in _strict_selected() and not _prefix_covered(code, _strict_external()) + ) + assert stranded == frozenset(), ( + f"the strict gate's RUF100 reads a valid `# noqa` for {sorted(stranded)} as unused, the " + "spurious-breach trap ruff-strict.toml's external override exists to prevent. Cover them " + "there." + ) + double_booked = frozenset( + code + for code in base_owned + if code in _strict_selected() and _prefix_covered(code, _strict_external()) + ) + assert double_booked == frozenset(), ( + f"{sorted(double_booked)} are selected by the strict config yet shadowed by its external " + "list, so their stale suppressions can never be reported. Narrow the external entry in " + "ruff-strict.toml." + ) + + +def test_every_budgeted_rule_is_one_the_gate_actually_measures(): + selectors = tuple(_lint_section("ruff-strict.toml")["select"]) + unmeasured = frozenset(code for code in _budgeted_rules() if not code.startswith(selectors)) + assert unmeasured == frozenset(), ( + f"the gate never counts {sorted(unmeasured)}, so their ceilings are dead config that reads " + "as coverage. Either select them in ruff-strict.toml or drop them from the budget." + ) + + +@_needs_ruff +def test_every_strict_selected_rule_is_budgeted_or_hard_failed_by_the_base_config(all_ruff_rule_codes): + strict_enabled = frozenset( + code + for code in all_ruff_rule_codes + if code.startswith(tuple(_lint_section("ruff-strict.toml")["select"])) + ) + base_hard_failed = tuple(_lint_section("ruff.toml")["extend-select"]) + unpoliced = frozenset( + code + for code in strict_enabled + if code not in _budgeted_rules() + and not code.startswith(base_hard_failed) + and code not in _ENABLED_BY_RUFF_DEFAULTS + ) + assert unpoliced == frozenset(), ( + f"nothing enforces {sorted(unpoliced)}: the gate skips rules missing from the budget, and " + "the base config does not hard-fail them. Re-add a budget ceiling or graduate them into " + "ruff.toml's lint.extend-select." + ) + + +@_needs_ruff +def test_a_noqa_for_a_strict_gate_rule_survives_the_normal_ruff_run(): + assert "RUF100" not in _ruff_output_for_noqa("ANN202") + + +@_needs_ruff +def test_the_external_list_is_what_saves_that_noqa(): + assert "RUF100" in _ruff_output_for_noqa("ANN202", "--config", "lint.external=[]") + + +@_needs_ruff +def test_a_stale_noqa_for_a_locally_enabled_rule_is_still_reported(): + assert "RUF100" in _ruff_output_for_noqa("F401") + + +def _ruff_output_for_source(source: str) -> str: + proc = subprocess.run( + [_RUFF, "check", "--no-cache", "--stdin-filename", "litellm/types/_graduate_probe.py", "-"], + cwd=_REPO_ROOT, + input=source, + capture_output=True, + text=True, + ) + return proc.stdout + + +_DEPRECATED_TYPING_ALIAS = "from typing import List # noqa: UP035\n\n\ndef _probe(x: List[int]) -> None: ...\n" + + +@_needs_ruff +def test_a_graduated_rule_now_fails_the_normal_ruff_run_instead_of_waiting_for_the_gate(): + assert "UP006" in _ruff_output_for_source(_DEPRECATED_TYPING_ALIAS) + + +@_needs_ruff +def test_a_graduated_rule_can_still_be_suppressed_without_tripping_unused_noqa(): + suppressed = _DEPRECATED_TYPING_ALIAS.replace("...\n", "... # noqa: UP006\n") + output = _ruff_output_for_source(suppressed) + assert "UP006" not in output + assert "RUF100" not in output From 5cd027cbbca7d731238bd870c01917e9dcf3af97 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:11:23 -0700 Subject: [PATCH 08/13] fix(lint): let the ratchet guard recognise a graduated rule A budget rule that graduates into a config's hard-fail select list rightly leaves the budget file, but the ratchet guard read any disappearance as a silently raised ceiling. Teach it the pairing between ruff-strict-budget.json and ruff.toml: a dropped rule is excused only when the paired config's lint.extend-select (minus lint.ignore) now hard-fails it, so deleting a rule without graduating it still trips the guard. --- scripts/budget_ratchet_check.py | 51 ++++++++++++++++--- .../test_litellm/test_budget_ratchet_check.py | 44 ++++++++++++++++ 2 files changed, 89 insertions(+), 6 deletions(-) diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 10a78483643..e97cd1bca00 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -10,7 +10,10 @@ content at the merge-base with the target branch and fails (exits 1, red) if: * a rule was dropped from a budget (its ceiling effectively became infinite), or * an entire budget file was deleted. -New rules and lowered/equal limits are fine. +New rules and lowered/equal limits are fine. So is a rule that graduated: once a +paired config (ruff.toml for the ruff-strict budget) selects the rule outright it +hard-fails at the first violation, which is stricter than any ceiling the budget +could hold, so dropping its entry tightens the guard rather than removing it. This is deliberately NOT a gating check. It should turn the run red so that a loosening is impossible to miss in review, but it must stay OUT of the @@ -30,7 +33,9 @@ import argparse import json import subprocess import sys +import tomllib from pathlib import Path +from types import MappingProxyType from typing import NamedTuple REPO_ROOT = Path(__file__).resolve().parent.parent @@ -40,6 +45,7 @@ DEFAULT_BUDGETS: tuple[str, ...] = ( "type-discipline-budget.json", "basedpyright-code-budget.json", ) +GRADUATION_CONFIGS = MappingProxyType({"ruff-strict-budget.json": "ruff.toml"}) class Regression(NamedTuple): @@ -106,24 +112,57 @@ def _limits(budget: dict) -> dict[str, int]: } +def selectors_hard_failed_by(lint: dict) -> tuple[str, ...]: + """A ruff `[lint]` table's selected codes, minus anything `ignore` turns back off. + + `lint.ignore` wins over `lint.extend-select` in ruff, so an ignored code is not + actually enforced and must not count as a graduation. + """ + ignored = tuple(lint.get("ignore", ())) + return tuple( + selector + for selector in lint.get("extend-select", ()) + if not (ignored and selector.startswith(ignored)) + ) + + +def graduated_selectors(rel: str) -> tuple[str, ...]: + """Selectors the budget's paired ruff config hard-fails, so its ceiling is moot.""" + config = GRADUATION_CONFIGS.get(rel) + if config is None or not (REPO_ROOT / config).exists(): + return () + return selectors_hard_failed_by( + tomllib.loads((REPO_ROOT / config).read_text()).get("lint", {}) + ) + + def _regression_detail( rule: str, base_limits: dict[str, int], head_limits: dict[str, int], + graduated: tuple[str, ...], ) -> str | None: - """Why `rule` regressed vs base, or None when it held flat or fell. + """Why `rule` regressed vs base, or None when it held flat, fell, or graduated. - A dropped rule is terminal; otherwise the only loosening left is a raised limit. + A dropped rule is terminal unless it graduated; otherwise the only loosening + left is a raised limit. """ base_limit = base_limits[rule] if rule not in head_limits: + if graduated and rule.startswith(graduated): + return None return f"rule dropped (limit {base_limit} -> removed)" if head_limits[rule] > base_limit: return f"limit raised {base_limit} -> {head_limits[rule]}" return None -def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]: +def regressions_for( + rel: str, + base: dict | None, + head: dict | None, + graduated: tuple[str, ...] = (), +) -> list[Regression]: if base is None: return [] # new budget file: nothing to ratchet against yet if head is None: @@ -133,7 +172,7 @@ def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regr return [ Regression(rel, rule, detail) for rule in sorted(base_limits) - if (detail := _regression_detail(rule, base_limits, head_limits)) is not None + if (detail := _regression_detail(rule, base_limits, head_limits, graduated)) is not None ] @@ -164,7 +203,7 @@ def main() -> int: print(f"skip {rel}: new file (no base at {args.base} to ratchet against)") continue checked.append(rel) - regressions.extend(regressions_for(rel, base, head)) + regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel))) if regressions: print( diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index 1972c1b6386..22d05f4d00d 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -68,6 +68,50 @@ def test_new_rule_in_head_is_clean(): assert ratchet.regressions_for("b.json", {}, {"new-rule": _spec_of(5)}) == [] +def test_dropped_rule_that_graduated_to_a_hard_failing_config_is_clean(): + base = {"UP006": _spec_of(0)} + assert ratchet.regressions_for("b.json", base, {}, graduated=("UP006",)) == [] + + +def test_graduation_matches_by_prefix_like_ruff_selectors_do(): + base = {"ANN202": _spec_of(865)} + assert ratchet.regressions_for("b.json", base, {}, graduated=("ANN",)) == [] + + +def test_an_unrelated_graduation_does_not_excuse_a_dropped_rule(): + base = {"C901": _spec_of(3)} + regs = ratchet.regressions_for("b.json", base, {}, graduated=("UP006", "SIM118")) + assert [r.rule for r in regs] == ["C901"] + assert "dropped" in regs[0].detail + + +def test_graduation_never_excuses_a_raised_limit(): + base = {"UP006": _spec_of(0)} + regs = ratchet.regressions_for("b.json", base, {"UP006": _spec_of(7)}, graduated=("UP006",)) + assert [r.rule for r in regs] == ["UP006"] + assert "0 -> 7" in regs[0].detail + + +def test_graduated_selectors_come_from_the_paired_ruff_config(): + selectors = ratchet.graduated_selectors("ruff-strict-budget.json") + assert "UP006" in selectors + assert "ANN" not in selectors + + +def test_budgets_without_a_paired_config_can_never_graduate(): + assert ratchet.graduated_selectors("type-discipline-budget.json") == () + assert ratchet.graduated_selectors("basedpyright-code-budget.json") == () + + +def test_a_selector_the_config_also_ignores_does_not_count_as_graduated(): + lint = {"ignore": ["UP006"], "extend-select": ["UP006", "SIM118"]} + assert ratchet.selectors_hard_failed_by(lint) == ("SIM118",) + + +def test_selectors_hard_failed_by_reads_a_config_with_no_ignore_list(): + assert ratchet.selectors_hard_failed_by({"extend-select": ["UP006"]}) == ("UP006",) + + def test_deleted_budget_file_is_a_regression(): regs = ratchet.regressions_for("b.json", {"LIT006": _spec_of(1)}, None) assert [r.rule for r in regs] == ["*"] From 5f7a663005bf3228f913d34cb372b712d318e7de Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:33:32 -0700 Subject: [PATCH 09/13] fix(proxy): enforce require_managed_files on every raw provider id route require_managed_files was only checked on upload, so raw provider ids still reached the batch, fine-tuning and vector store file routes. Ownership rows exist only for managed ids, so those requests were forwarded under shared credentials with no tenant check: knowing another tenant's id was enough to read, run against, cancel or delete their object. Generalise the file-id guard to validate_managed_id_requirement(resource_id, resource_kind) and call it on batch create/retrieve/cancel, fine-tuning create/retrieve/cancel (training_file and validation_file both) and the shared vector store file id resolver. Behaviour is unchanged when the setting is off. --- litellm/proxy/batches_endpoints/endpoints.py | 5 + .../proxy/fine_tuning_endpoints/endpoints.py | 8 + .../openai_files_endpoints/common_utils.py | 27 +- .../openai_files_endpoints/files_endpoints.py | 8 +- .../vector_store_files_endpoints/endpoints.py | 5 + .../proxy/batches_endpoints/test_endpoints.py | 117 +++++++++ .../proxy/fine_tuning_endpoints/__init__.py | 0 .../fine_tuning_endpoints/test_endpoints.py | 236 ++++++++++++++++++ .../test_files_endpoint.py | 8 +- .../vector_store_files_endpoints/__init__.py | 0 .../test_endpoints.py | 83 ++++++ 11 files changed, 479 insertions(+), 18 deletions(-) create mode 100644 tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py create mode 100644 tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index f7c332f2849..c9f66c5a48b 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_original_file_id, prepare_data_with_credentials, update_batch_in_database, + validate_managed_id_requirement, ) from litellm.proxy.utils import handle_exception_on_proxy, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository @@ -176,6 +177,7 @@ async def create_batch( } input_file_id: Final = _create_batch_data.get("input_file_id", None) + validate_managed_id_requirement(resource_id=input_file_id, resource_kind="file") unified_file_id: str | Literal[False] = False model_from_file_id = None @@ -392,6 +394,7 @@ async def retrieve_batch( data: dict = {} try: + validate_managed_id_requirement(resource_id=batch_id, resource_kind="batch") model_from_id: Final = decode_model_from_file_id(batch_id) _retrieve_batch_request: Final = RetrieveBatchRequest( batch_id=batch_id, @@ -840,6 +843,8 @@ async def cancel_batch( data: dict = {} try: + validate_managed_id_requirement(resource_id=batch_id, resource_kind="batch") + # Check for encoded batch ID with model info model_from_id: Final = decode_model_from_file_id(batch_id) diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index f8ffb77edb8..87e9895eed0 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + validate_managed_id_requirement, ) from litellm.proxy.utils import handle_exception_on_proxy from litellm.types.utils import LiteLLMFineTuningJob @@ -134,6 +135,11 @@ async def create_fine_tuning_job( ## CHECK IF MANAGED FILE ID unified_file_id: str | Literal[False] = False training_file: Final = fine_tuning_request.training_file + validate_managed_id_requirement(resource_id=training_file, resource_kind="file") + validate_managed_id_requirement( + resource_id=fine_tuning_request.validation_file, + resource_kind="file", + ) response: LiteLLMFineTuningJob | None = None if training_file: unified_file_id = _is_base64_encoded_unified_file_id(training_file) @@ -246,6 +252,7 @@ async def retrieve_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") + validate_managed_id_requirement(resource_id=fine_tuning_job_id, resource_kind="fine-tuning job") # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( @@ -513,6 +520,7 @@ async def cancel_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") + validate_managed_id_requirement(resource_id=fine_tuning_job_id, resource_kind="fine-tuning job") # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b26010c5597..143bd5bc6b5 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -881,17 +881,20 @@ def validate_managed_files_requirement( ) -def validate_managed_file_id_requirement(file_id: str) -> None: +def validate_managed_id_requirement( + resource_id: str | None, + resource_kind: Literal["file", "batch", "fine-tuning job"], +) -> None: """ - Enforce proxy-level managed files on the file read/delete routes when - ``litellm.require_managed_files`` is enabled. + Enforce proxy-level managed resources on every route that accepts a provider-issued id + when ``litellm.require_managed_files`` is enabled. - Ownership is only recorded for LiteLLM managed files, so a raw provider file id sent to - retrieve/content/delete is forwarded to the provider under shared credentials without any - tenant check; knowing another tenant's provider file id would be enough to read or delete it. + Ownership is only recorded for LiteLLM managed ids, so a raw provider id is forwarded to the + provider under shared credentials without any tenant check; knowing another tenant's provider + id would be enough to read, reuse, or destroy the object behind it. Raises: - HTTPException: 400 if ``file_id`` is not a LiteLLM managed file id. + HTTPException: 400 if ``resource_id`` is set and is not a LiteLLM managed id. """ from fastapi import HTTPException @@ -900,14 +903,18 @@ def validate_managed_file_id_requirement(file_id: str) -> None: if litellm.require_managed_files is not True: return - if _is_base64_encoded_unified_file_id(file_id): + if not resource_id: + return + + if _is_base64_encoded_unified_file_id(resource_id): return raise HTTPException( status_code=400, detail=( - "Raw provider file ids cannot be used when require_managed_files is enabled in " - "litellm_settings. Use the LiteLLM managed file id returned when the file was created." + f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in " + f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the " + f"{resource_kind} was created." ), ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 5c1a4441d0d..b5950c4f852 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -49,8 +49,8 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, - validate_managed_file_id_requirement, validate_managed_files_requirement, + validate_managed_id_requirement, ) from litellm.proxy.utils import ProxyLogging, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository @@ -613,7 +613,7 @@ async def get_file_content( data: dict = {"file_id": file_id} try: - validate_managed_file_id_requirement(file_id=file_id) + validate_managed_id_requirement(resource_id=file_id, resource_kind="file") # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) @@ -911,7 +911,7 @@ async def get_file( data: dict = {"file_id": file_id} try: - validate_managed_file_id_requirement(file_id=file_id) + validate_managed_id_requirement(resource_id=file_id, resource_kind="file") custom_llm_provider: Final = ( provider @@ -1103,7 +1103,7 @@ async def delete_file( data: dict = {"file_id": file_id} try: - validate_managed_file_id_requirement(file_id=file_id) + validate_managed_id_requirement(resource_id=file_id, resource_kind="file") custom_llm_provider: Final = ( provider diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 896b7ca33d7..f8fdf607292 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -65,6 +65,11 @@ def _update_request_data_with_managed_file_id( is_base64_encoded_unified_id, parse_unified_id, ) + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_id_requirement, + ) + + validate_managed_id_requirement(resource_id=file_id, resource_kind="file") # First, check if this is a unified managed file ID (base64 encoded) decoded_id: Final = is_base64_encoded_unified_id(file_id) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index e758aa5ca7f..18ca604f88d 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -2253,3 +2253,120 @@ async def test_cancel__provider_only_resolves_named_vertex_credentials(cancel_ha "vertex_location": "us-central1", "vertex_credentials": "/creds/customer-sa.json", } + + +# =========================================================================== # +# require_managed_files - raw provider ids must not reach the provider. # +# # +# Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to # +# these routes is forwarded under the shared provider credentials with no # +# tenant check, so any caller who learns another tenant's id can read its # +# batch, reuse its file as batch input, or cancel its job. These lock the # +# guard on every batches route that accepts a caller-supplied id. # +# =========================================================================== # + + +def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-provider-id") -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +@pytest.mark.asyncio +async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness): + set_body( + harness, + { + "input_file_id": "file-victim-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "400" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__model_encoded_input_file_id_rejected_when_managed_files_required(harness): + """A model-encoded id is client-forgeable and has no ownership row, so it is + not a managed file id and must be rejected like any other raw id.""" + set_body( + harness, + { + "input_file_id": AZURE_FILE_ID, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "400" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__raw_input_file_id_allowed_when_managed_files_not_required(harness): + set_body( + harness, + { + "input_file_id": "file-victim-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", False): + await call_create(harness) + + assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123" + + +@pytest.mark.asyncio +async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_retrieve(retrieve_harness, "batch-victim-abc123") + + assert exc.value.code == "400" + retrieve_harness.litellm_aretrieve.assert_not_called() + retrieve_harness.router_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness): + with patch.object(litellm, "require_managed_files", True): + await call_retrieve(retrieve_harness, _unified_batch_id()) + + assert retrieve_harness.router_aretrieve.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_harness): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_cancel(cancel_harness, "batch-victim-abc123") + + assert exc.value.code == "400" + cancel_harness.litellm_acancel.assert_not_called() + cancel_harness.router_acancel.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness): + with patch.object(litellm, "require_managed_files", True): + await call_cancel(cancel_harness, _unified_batch_id()) + + assert cancel_harness.router_acancel.call_count == 1 diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py b/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py new file mode 100644 index 00000000000..35c202057e7 --- /dev/null +++ b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py @@ -0,0 +1,236 @@ +""" +require_managed_files enforcement for litellm/proxy/fine_tuning_endpoints/endpoints.py + +Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to these +routes is forwarded to the provider under the shared proxy credentials with no tenant +check, so any caller who learns another tenant's file id can train on it, and any +caller who learns another tenant's job id can read or cancel it. + +Each test asserts BOTH that the request is rejected AND that every downstream provider +seam stayed untouched, so a guard that raises after the provider call would still fail. +""" + +import base64 +import os +import sys +from contextlib import ExitStack +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import Response + +import litellm +import litellm.proxy.fine_tuning_endpoints.endpoints as endpoints +import litellm.proxy.proxy_server as proxy_server +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.types.llms.openai import LiteLLMFineTuningJobCreate +from litellm.types.utils import LiteLLMFineTuningJob, SpecialEnums + +RAW_FILE_ID = "file-victim-abc123" +RAW_JOB_ID = "ftjob-victim-abc123" + + +def _unified_file_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _unified_job_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format("gpt-4o-mini-id", RAW_JOB_ID) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _job() -> LiteLLMFineTuningJob: + job = LiteLLMFineTuningJob( + id=RAW_JOB_ID, + created_at=1234567890, + fine_tuned_model=None, + finished_at=None, + hyperparameters={"n_epochs": 1}, + model="gpt-4o-mini", + object="fine_tuning.job", + organization_id="org-test", + result_files=[], + seed=0, + status="running", + trained_tokens=None, + training_file=RAW_FILE_ID, + validation_file=None, + ) + job._hidden_params = {} + return job + + +class FakeRequest: + def __init__(self): + self.headers = {} + self.query_params = {} + + async def json(self): + return {} + + +class Seams: + def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock]): + self.router = router + self.litellm_calls = litellm_calls + + def assert_no_provider_call(self) -> None: + for name, mock in self.litellm_calls.items(): + assert mock.call_count == 0, f"litellm.{name} was called" + for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job"): + assert getattr(self.router, name).call_count == 0, f"router.{name} was called" + + +@pytest.fixture +def seams(): + logging = MagicMock(spec=ProxyLogging) + logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"]) + logging.post_call_failure_hook = AsyncMock() + logging.update_request_status = AsyncMock() + logging.get_proxy_hook = MagicMock(return_value=None) + + router = MagicMock(spec=Router) + router.acreate_fine_tuning_job = AsyncMock(return_value=_job()) + router.aretrieve_fine_tuning_job = AsyncMock(return_value=_job()) + router.acancel_fine_tuning_job = AsyncMock(return_value=_job()) + + litellm_calls = { + name: AsyncMock(return_value=_job()) + for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job") + } + + with ExitStack() as stack: + stack.enter_context( + patch.object( + ProxyBaseLLMRequestProcessing, + "common_processing_pre_call_logic", + AsyncMock(side_effect=lambda self=None, **kw: (self.data if self else {}, MagicMock())), + ) + ) + stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", MagicMock(return_value={}))) + for name, mock in litellm_calls.items(): + stack.enter_context(patch.object(litellm, name, mock)) + stack.enter_context(patch.object(proxy_server, "llm_router", router)) + stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging)) + stack.enter_context(patch.object(proxy_server, "premium_user", True)) + stack.enter_context(patch.object(proxy_server, "general_settings", {})) + stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock())) + stack.enter_context(patch.object(proxy_server, "version", "test-version")) + stack.enter_context(patch.object(endpoints, "fine_tuning_config", [{"custom_llm_provider": "openai"}])) + yield Seams(router=router, litellm_calls=litellm_calls) + + +async def _create(training_file: str, validation_file: str | None = None): + return await endpoints.create_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_request=LiteLLMFineTuningJobCreate( + model="gpt-4o-mini", + training_file=training_file, + validation_file=validation_file, + custom_llm_provider="openai", + ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +async def _retrieve(job_id: str): + return await endpoints.retrieve_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_job_id=job_id, + custom_llm_provider="openai", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +async def _cancel(job_id: str): + return await endpoints.cancel_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_job_id=job_id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +@pytest.mark.asyncio +async def test_create__raw_training_file_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(RAW_FILE_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__raw_validation_file_rejected_when_managed_files_required(seams): + """The validation file is uploaded and readable exactly like the training file, + so a managed training_file must not smuggle a raw validation_file past the guard.""" + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(_unified_file_id(), validation_file=RAW_FILE_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__unified_training_file_allowed_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + await _create(_unified_file_id()) + + assert seams.router.acreate_fine_tuning_job.call_count == 1 + + +@pytest.mark.asyncio +async def test_create__raw_training_file_allowed_when_managed_files_not_required(seams): + with patch.object(litellm, "require_managed_files", False): + await _create(RAW_FILE_ID) + + assert seams.litellm_calls["acreate_fine_tuning_job"].call_count == 1 + + +@pytest.mark.asyncio +async def test_retrieve__raw_job_id_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _retrieve(RAW_JOB_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_retrieve__unified_job_id_allowed_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + await _retrieve(_unified_job_id()) + + assert seams.router.aretrieve_fine_tuning_job.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancel__raw_job_id_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _cancel(RAW_JOB_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_cancel__unified_job_id_allowed_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + await _cancel(_unified_job_id()) + + assert seams.router.acancel_fine_tuning_job.call_count == 1 diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 24b814bae1f..2e15da7590a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3109,22 +3109,22 @@ def _unified_managed_file_id() -> str: def test_require_managed_files_allows_unified_managed_file_id(monkeypatch): from litellm.proxy.openai_files_endpoints.common_utils import ( - validate_managed_file_id_requirement, + validate_managed_id_requirement, ) monkeypatch.setattr("litellm.require_managed_files", True) - validate_managed_file_id_requirement(file_id=_unified_managed_file_id()) + validate_managed_id_requirement(resource_id=_unified_managed_file_id(), resource_kind="file") def test_managed_file_id_requirement_is_opt_in(monkeypatch): from litellm.proxy.openai_files_endpoints.common_utils import ( - validate_managed_file_id_requirement, + validate_managed_id_requirement, ) monkeypatch.setattr("litellm.require_managed_files", False) - validate_managed_file_id_requirement(file_id="file-victim-abc123") + validate_managed_id_requirement(resource_id="file-victim-abc123", resource_kind="file") def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py b/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py new file mode 100644 index 00000000000..c2bdb1d80f2 --- /dev/null +++ b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py @@ -0,0 +1,83 @@ +""" +require_managed_files enforcement for litellm/proxy/vector_store_files_endpoints/endpoints.py + +Every vector-store file route (create, retrieve, content, update, delete) resolves its +caller-supplied file id through _update_request_data_with_managed_file_id before the +provider call, so the guard lives there once and covers all five. + +A raw provider file id has no ownership row, so without the guard it is attached to a +vector store or read back under the shared provider credentials with no tenant check. +""" + +import base64 +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import HTTPException + +import litellm +from litellm.proxy.vector_store_files_endpoints.endpoints import ( + _update_request_data_with_managed_file_id, +) +from litellm.types.utils import SpecialEnums + +RAW_FILE_ID = "file-victim-abc123" + + +def _unified_file_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _resolve(file_id: str): + return _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": file_id}, + file_id=file_id, + request=MagicMock(headers={}, query_params={}), + llm_router=None, + ) + + +def test_raw_file_id_rejected_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + _resolve(RAW_FILE_ID) + + assert exc.value.status_code == 400 + + +def test_model_encoded_file_id_rejected_when_managed_files_required(): + """encode_file_id_with_model output is client-forgeable and carries no ownership + row, so it is not a managed file id.""" + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + encoded = encode_file_id_with_model(RAW_FILE_ID, "gpt-4o-mini", id_type="file") + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + _resolve(encoded) + + assert exc.value.status_code == 400 + + +def test_unified_file_id_allowed_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + data, original = _resolve(_unified_file_id()) + + assert original == _unified_file_id() + assert data["file_id"] == RAW_FILE_ID + + +def test_raw_file_id_allowed_when_managed_files_not_required(): + with patch.object(litellm, "require_managed_files", False): + data, original = _resolve(RAW_FILE_ID) + + assert original is None + assert data["file_id"] == RAW_FILE_ID From b01eacd67c3c8828c464f29ff5fbaee32f0735ce Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:36:30 -0700 Subject: [PATCH 10/13] ci: run the new fine-tuning and vector store file test dirs --- .github/workflows/test-unit-proxy-endpoints.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index 645996f779d..2ea3c521e8b 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -38,6 +38,8 @@ jobs: tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/batches_endpoints + tests/test_litellm/proxy/fine_tuning_endpoints + tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/image_endpoints From 8c0556abf6965fde0de260da5ce424aa1daa1a56 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 02:42:34 -0700 Subject: [PATCH 11/13] fix(proxy): authenticate managed ids before routing --- litellm/proxy/batches_endpoints/endpoints.py | 21 ++++- .../proxy/fine_tuning_endpoints/endpoints.py | 25 +++++- .../openai_files_endpoints/common_utils.py | 58 ++++++++++--- .../openai_files_endpoints/files_endpoints.py | 21 ++++- .../vector_store_files_endpoints/endpoints.py | 56 +++++++++--- .../proxy/batches_endpoints/test_endpoints.py | 56 ++++++++++++ .../fine_tuning_endpoints/test_endpoints.py | 46 +++++++++- .../test_files_endpoint.py | 36 +++++++- .../test_endpoints.py | 87 ++++++++++++++++--- 9 files changed, 355 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index c9f66c5a48b..aef1c5ac17e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -177,7 +177,12 @@ async def create_batch( } input_file_id: Final = _create_batch_data.get("input_file_id", None) - validate_managed_id_requirement(resource_id=input_file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=input_file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) unified_file_id: str | Literal[False] = False model_from_file_id = None @@ -394,7 +399,12 @@ async def retrieve_batch( data: dict = {} try: - validate_managed_id_requirement(resource_id=batch_id, resource_kind="batch") + await validate_managed_id_requirement( + resource_id=batch_id, + resource_kind="batch", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) model_from_id: Final = decode_model_from_file_id(batch_id) _retrieve_batch_request: Final = RetrieveBatchRequest( batch_id=batch_id, @@ -843,7 +853,12 @@ async def cancel_batch( data: dict = {} try: - validate_managed_id_requirement(resource_id=batch_id, resource_kind="batch") + await validate_managed_id_requirement( + resource_id=batch_id, + resource_kind="batch", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Check for encoded batch ID with model info model_from_id: Final = decode_model_from_file_id(batch_id) diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index 87e9895eed0..a13ad00713d 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -135,10 +135,17 @@ async def create_fine_tuning_job( ## CHECK IF MANAGED FILE ID unified_file_id: str | Literal[False] = False training_file: Final = fine_tuning_request.training_file - validate_managed_id_requirement(resource_id=training_file, resource_kind="file") - validate_managed_id_requirement( + await validate_managed_id_requirement( + resource_id=training_file, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + await validate_managed_id_requirement( resource_id=fine_tuning_request.validation_file, resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), ) response: LiteLLMFineTuningJob | None = None if training_file: @@ -252,7 +259,12 @@ async def retrieve_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") - validate_managed_id_requirement(resource_id=fine_tuning_job_id, resource_kind="fine-tuning job") + await validate_managed_id_requirement( + resource_id=fine_tuning_job_id, + resource_kind="fine-tuning job", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( @@ -520,7 +532,12 @@ async def cancel_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") - validate_managed_id_requirement(resource_id=fine_tuning_job_id, resource_kind="fine-tuning job") + await validate_managed_id_requirement( + resource_id=fine_tuning_job_id, + resource_kind="fine-tuning job", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 143bd5bc6b5..56e986c89cf 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -4,7 +4,7 @@ import re from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, runtime_checkable from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -22,6 +22,21 @@ if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch +@runtime_checkable +class ManagedResourceAccessChecker(Protocol): + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> bool: ... + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> bool: ... + + def _is_base64_encoded_unified_file_id(b64_uid: str) -> str | Literal[False]: # Ensure b64_uid is a string and not a mock object if not isinstance(b64_uid, str): @@ -881,20 +896,24 @@ def validate_managed_files_requirement( ) -def validate_managed_id_requirement( +async def validate_managed_id_requirement( resource_id: str | None, resource_kind: Literal["file", "batch", "fine-tuning job"], + user_api_key_dict: "UserAPIKeyAuth", + managed_files_obj: object | None, ) -> None: """ Enforce proxy-level managed resources on every route that accepts a provider-issued id - when ``litellm.require_managed_files`` is enabled. + when ``litellm.require_managed_files`` is enabled, and authenticate managed ids against + the caller's stored ownership record. Ownership is only recorded for LiteLLM managed ids, so a raw provider id is forwarded to the provider under shared credentials without any tenant check; knowing another tenant's provider id would be enough to read, reuse, or destroy the object behind it. Raises: - HTTPException: 400 if ``resource_id`` is set and is not a LiteLLM managed id. + HTTPException: 400 for a raw id, 403 for an inaccessible managed id, or 500 when + ownership validation is unavailable. """ from fastapi import HTTPException @@ -906,16 +925,33 @@ def validate_managed_id_requirement( if not resource_id: return - if _is_base64_encoded_unified_file_id(resource_id): + if not _is_base64_encoded_unified_file_id(resource_id): + raise HTTPException( + status_code=400, + detail=( + f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in " + f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the " + f"{resource_kind} was created." + ), + ) + + if not isinstance(managed_files_obj, ManagedResourceAccessChecker): + raise HTTPException( + status_code=500, + detail="Managed resource ownership validation is unavailable.", + ) + + can_access: Final = ( + await managed_files_obj.can_user_call_unified_file_id(resource_id, user_api_key_dict) + if resource_kind == "file" + else await managed_files_obj.can_user_call_unified_object_id(resource_id, user_api_key_dict) + ) + if can_access: return raise HTTPException( - status_code=400, - detail=( - f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in " - f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the " - f"{resource_kind} was created." - ), + status_code=403, + detail=f"The caller does not have access to this managed {resource_kind} id.", ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index b5950c4f852..0acaac3bf5d 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -613,7 +613,12 @@ async def get_file_content( data: dict = {"file_id": file_id} try: - validate_managed_id_requirement(resource_id=file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) @@ -911,7 +916,12 @@ async def get_file( data: dict = {"file_id": file_id} try: - validate_managed_id_requirement(resource_id=file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) custom_llm_provider: Final = ( provider @@ -1103,7 +1113,12 @@ async def delete_file( data: dict = {"file_id": file_id} try: - validate_managed_id_requirement(resource_id=file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) custom_llm_provider: Final = ( provider diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index f8fdf607292..c9b89bcd390 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -30,10 +30,12 @@ if TYPE_CHECKING: router: Final = APIRouter() -def _update_request_data_with_managed_file_id( +async def _update_request_data_with_managed_file_id( data: dict, file_id: str, request: Request, + user_api_key_dict: UserAPIKeyAuth, + managed_files_obj: object | None, llm_router: Optional["Router"] = None, ) -> tuple[dict, str | None]: """ @@ -69,7 +71,12 @@ def _update_request_data_with_managed_file_id( validate_managed_id_requirement, ) - validate_managed_id_requirement(resource_id=file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=managed_files_obj, + ) # First, check if this is a unified managed file ID (base64 encoded) decoded_id: Final = is_base64_encoded_unified_id(file_id) @@ -514,8 +521,13 @@ async def vector_store_file_create( # Handle managed file IDs if present in request body original_managed_file_id = None if "file_id" in data: - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=data["file_id"], request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=data["file_id"], + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -712,8 +724,13 @@ async def vector_store_file_retrieve( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -814,8 +831,13 @@ async def vector_store_file_content( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -916,8 +938,13 @@ async def vector_store_file_update( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -1018,8 +1045,13 @@ async def vector_store_file_delete( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 18ca604f88d..f9193db143e 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -2275,6 +2275,37 @@ def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-pro return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") +def _unified_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "managed-id", "gpt-4o-mini", "file-provider-id", "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: bool = True + object_access: bool = True + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.file_access + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.object_access + + @pytest.mark.asyncio async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness): set_body( @@ -2334,6 +2365,27 @@ async def test_create__raw_input_file_id_allowed_when_managed_files_not_required assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123" +@pytest.mark.asyncio +async def test_create__other_teams_unified_input_file_id_rejected(harness): + set_body( + harness, + { + "input_file_id": _unified_file_id(), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "403" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + @pytest.mark.asyncio async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness): with patch.object(litellm, "require_managed_files", True): @@ -2347,6 +2399,8 @@ async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retri @pytest.mark.asyncio async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness): + retrieve_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await call_retrieve(retrieve_harness, _unified_batch_id()) @@ -2366,6 +2420,8 @@ async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_ @pytest.mark.asyncio async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness): + cancel_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await call_cancel(cancel_harness, _unified_batch_id()) diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py index 35c202057e7..b54787bf428 100644 --- a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py @@ -14,6 +14,7 @@ import base64 import os import sys from contextlib import ExitStack +from dataclasses import dataclass from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -78,10 +79,31 @@ class FakeRequest: return {} +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: bool = True + object_access: bool = True + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.file_access + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.object_access + + class Seams: - def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock]): + def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock], logging: MagicMock): self.router = router self.litellm_calls = litellm_calls + self.logging = logging def assert_no_provider_call(self) -> None: for name, mock in self.litellm_calls.items(): @@ -126,7 +148,7 @@ def seams(): stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock())) stack.enter_context(patch.object(proxy_server, "version", "test-version")) stack.enter_context(patch.object(endpoints, "fine_tuning_config", [{"custom_llm_provider": "openai"}])) - yield Seams(router=router, litellm_calls=litellm_calls) + yield Seams(router=router, litellm_calls=litellm_calls, logging=logging) async def _create(training_file: str, validation_file: str | None = None): @@ -176,6 +198,8 @@ async def test_create__raw_training_file_rejected_when_managed_files_required(se async def test_create__raw_validation_file_rejected_when_managed_files_required(seams): """The validation file is uploaded and readable exactly like the training file, so a managed training_file must not smuggle a raw validation_file past the guard.""" + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): with pytest.raises(ProxyException) as exc: await _create(_unified_file_id(), validation_file=RAW_FILE_ID) @@ -186,12 +210,26 @@ async def test_create__raw_validation_file_rejected_when_managed_files_required( @pytest.mark.asyncio async def test_create__unified_training_file_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await _create(_unified_file_id()) assert seams.router.acreate_fine_tuning_job.call_count == 1 +@pytest.mark.asyncio +async def test_create__other_teams_unified_training_file_rejected(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(_unified_file_id()) + + assert exc.value.code == "403" + seams.assert_no_provider_call() + + @pytest.mark.asyncio async def test_create__raw_training_file_allowed_when_managed_files_not_required(seams): with patch.object(litellm, "require_managed_files", False): @@ -212,6 +250,8 @@ async def test_retrieve__raw_job_id_rejected_when_managed_files_required(seams): @pytest.mark.asyncio async def test_retrieve__unified_job_id_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await _retrieve(_unified_job_id()) @@ -230,6 +270,8 @@ async def test_cancel__raw_job_id_rejected_when_managed_files_required(seams): @pytest.mark.asyncio async def test_cancel__unified_job_id_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await _cancel(_unified_job_id()) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 2e15da7590a..e68e7102fce 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3107,24 +3107,52 @@ def _unified_managed_file_id() -> str: return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") -def test_require_managed_files_allows_unified_managed_file_id(monkeypatch): +class _ManagedResourceAccessCheckerStub: + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return True + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return True + + +@pytest.mark.asyncio +async def test_require_managed_files_allows_owned_unified_managed_file_id(monkeypatch): from litellm.proxy.openai_files_endpoints.common_utils import ( validate_managed_id_requirement, ) monkeypatch.setattr("litellm.require_managed_files", True) - validate_managed_id_requirement(resource_id=_unified_managed_file_id(), resource_kind="file") + await validate_managed_id_requirement( + resource_id=_unified_managed_file_id(), + resource_kind="file", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="owner-user"), + managed_files_obj=_ManagedResourceAccessCheckerStub(), + ) -def test_managed_file_id_requirement_is_opt_in(monkeypatch): +@pytest.mark.asyncio +async def test_managed_file_id_requirement_is_opt_in(monkeypatch): from litellm.proxy.openai_files_endpoints.common_utils import ( validate_managed_id_requirement, ) monkeypatch.setattr("litellm.require_managed_files", False) - validate_managed_id_requirement(resource_id="file-victim-abc123", resource_kind="file") + await validate_managed_id_requirement( + resource_id="file-victim-abc123", + resource_kind="file", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + managed_files_obj=None, + ) def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py index c2bdb1d80f2..da5dd1934e4 100644 --- a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py @@ -5,13 +5,15 @@ Every vector-store file route (create, retrieve, content, update, delete) resolv caller-supplied file id through _update_request_data_with_managed_file_id before the provider call, so the guard lives there once and covers all five. -A raw provider file id has no ownership row, so without the guard it is attached to a -vector store or read back under the shared provider credentials with no tenant check. +A raw or forged managed-looking file id has no ownership row, so without the guard it +is attached to a vector store or read back under shared provider credentials. """ import base64 import os import sys +from dataclasses import dataclass +from typing import Literal from unittest.mock import MagicMock, patch import pytest @@ -21,12 +23,35 @@ sys.path.insert(0, os.path.abspath("../../../..")) from fastapi import HTTPException import litellm +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.vector_store_files_endpoints.endpoints import ( _update_request_data_with_managed_file_id, ) from litellm.types.utils import SpecialEnums RAW_FILE_ID = "file-victim-abc123" +CALLER = UserAPIKeyAuth(api_key="sk-test", user_id="attacker-user", team_id="team-b") + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: Literal["allow", "deny", "missing"] + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + if self.file_access == "missing": + raise HTTPException(status_code=404, detail=f"File not found: {unified_file_id}") + return self.file_access == "allow" + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return False def _unified_file_id() -> str: @@ -36,24 +61,31 @@ def _unified_file_id() -> str: return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") -def _resolve(file_id: str): - return _update_request_data_with_managed_file_id( +async def _resolve( + file_id: str, + file_access: Literal["allow", "deny", "missing"] = "allow", +): + return await _update_request_data_with_managed_file_id( data={"vector_store_id": "vs-test", "file_id": file_id}, file_id=file_id, request=MagicMock(headers={}, query_params={}), + user_api_key_dict=CALLER, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access=file_access), llm_router=None, ) -def test_raw_file_id_rejected_when_managed_files_required(): +@pytest.mark.asyncio +async def test_raw_file_id_rejected_when_managed_files_required(): with patch.object(litellm, "require_managed_files", True): with pytest.raises(HTTPException) as exc: - _resolve(RAW_FILE_ID) + await _resolve(RAW_FILE_ID) assert exc.value.status_code == 400 -def test_model_encoded_file_id_rejected_when_managed_files_required(): +@pytest.mark.asyncio +async def test_model_encoded_file_id_rejected_when_managed_files_required(): """encode_file_id_with_model output is client-forgeable and carries no ownership row, so it is not a managed file id.""" from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model @@ -62,22 +94,53 @@ def test_model_encoded_file_id_rejected_when_managed_files_required(): with patch.object(litellm, "require_managed_files", True): with pytest.raises(HTTPException) as exc: - _resolve(encoded) + await _resolve(encoded) assert exc.value.status_code == 400 -def test_unified_file_id_allowed_when_managed_files_required(): +@pytest.mark.asyncio +async def test_forged_unified_file_id_rejected_without_ownership_record(): + forged_id = _unified_file_id() + data = {"vector_store_id": "vs-test", "file_id": forged_id} + with patch.object(litellm, "require_managed_files", True): - data, original = _resolve(_unified_file_id()) + with pytest.raises(HTTPException) as exc: + await _update_request_data_with_managed_file_id( + data=data, + file_id=forged_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=CALLER, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access="missing"), + llm_router=None, + ) + + assert exc.value.status_code == 404 + assert data["file_id"] == forged_id + + +@pytest.mark.asyncio +async def test_other_teams_unified_file_id_rejected(): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _resolve(_unified_file_id(), file_access="deny") + + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_owned_unified_file_id_allowed_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + data, original = await _resolve(_unified_file_id()) assert original == _unified_file_id() assert data["file_id"] == RAW_FILE_ID -def test_raw_file_id_allowed_when_managed_files_not_required(): +@pytest.mark.asyncio +async def test_raw_file_id_allowed_when_managed_files_not_required(): with patch.object(litellm, "require_managed_files", False): - data, original = _resolve(RAW_FILE_ID) + data, original = await _resolve(RAW_FILE_ID) assert original is None assert data["file_id"] == RAW_FILE_ID From 20eb7bb43718958bfac0e06225ead0b5ddb1d5b7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 08:14:29 -0700 Subject: [PATCH 12/13] chore(typing): clear 1.4k basedpyright Any errors across 21 hotspot files Typing-only pass over the 21 files with the highest reportAny and reportExplicitAny density among self-contained modules: management endpoints, guardrails, streaming internals, response transformations, MCP server, enterprise managed files, and vector store management. Whole-tree basedpyright drops from 148,648 to 146,984 errors (-1,664), with reportAny -1,111 and reportExplicitAny -296. No rule increased repo-wide and no file regressed on any rule. No cast(), type: ignore, noqa, suppression comments, or new Any annotations anywhere in the diff, and no runtime behavior changes. Budgets ratcheted by make lint-budget-update: basedpyright -1,663 across 48 rules, ruff-strict -86, type-discipline -110. --- basedpyright-code-budget.json | 28 +- .../proxy/hooks/managed_files.py | 524 +++++++----------- .../pydantic_ai_agents/transformation.py | 265 +++++---- .../websearch_interception/handler.py | 12 +- .../litellm_core_utils/realtime_streaming.py | 55 +- .../streaming_chunk_builder_utils.py | 97 +++- .../adapters/handler.py | 84 +-- litellm/llms/azure/assistants.py | 36 +- .../mcp_server/rest_endpoints.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 91 +-- .../proxy/guardrails/guardrail_endpoints.py | 173 ++++-- .../cisco_ai_defense/cisco_ai_defense.py | 100 ++-- .../unified_guardrail/unified_guardrail.py | 136 +++-- litellm/proxy/hooks/litellm_skills/main.py | 60 +- .../key_management_endpoints.py | 77 +-- .../model_management_endpoints.py | 160 ++++-- .../management_endpoints/team_endpoints.py | 224 +++++--- litellm/proxy/management_endpoints/ui_sso.py | 187 +++++-- .../proxy_setting_endpoints.py | 123 +++- .../management_endpoints.py | 72 ++- .../transformation.py | 184 +++--- litellm/responses/streaming_iterator.py | 51 +- ruff-strict-budget.json | 10 +- type-discipline-budget.json | 8 +- 24 files changed, 1644 insertions(+), 1115 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 32b8eb3d4d0..0385f7a96e7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 28842 + "limit": 27731 }, "reportArgumentType": { - "limit": 2634 + "limit": 2626 }, "reportAssignmentType": { "limit": 329 @@ -12,7 +12,7 @@ "limit": 514 }, "reportCallIssue": { - "limit": 117 + "limit": 116 }, "reportConstantRedefinition": { "limit": 40 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9103 + "limit": 8807 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5843 + "limit": 5835 }, "reportMissingTypeArgument": { - "limit": 15816 + "limit": 15790 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1078 + "limit": 1077 }, "reportOptionalOperand": { "limit": 0 @@ -90,28 +90,28 @@ "limit": 8 }, "reportReturnType": { - "limit": 218 + "limit": 217 }, "reportTypedDictNotRequiredAccess": { - "limit": 27 + "limit": 26 }, "reportUndefinedVariable": { "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45098 + "limit": 45063 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39826 + "limit": 39773 }, "reportUnknownParameterType": { - "limit": 20237 + "limit": 20207 }, "reportUnknownVariableType": { - "limit": 31371 + "limit": 31281 }, "reportUnnecessaryCast": { "limit": 122 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 864 + "limit": 862 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index d8318962633..f0914240f79 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,9 +3,21 @@ import base64 import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Final, + List, + Literal, + Optional, + Protocol, + TypedDict, + Union, + cast, +) from uuid import NAMESPACE_URL, uuid5 from fastapi import HTTPException @@ -98,33 +110,76 @@ def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMB try: batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object)) except Exception as e: - verbose_logger.warning( - f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}" - ) + verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}") return None batch_obj.id = row.unified_object_id return batch_obj -def _parse_managed_file_object( - raw_file_object: object, unified_file_id: str -) -> Optional[OpenAIFileObject]: +def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> Optional[OpenAIFileObject]: if raw_file_object is None: return None try: return OpenAIFileObject.model_validate(raw_file_object) except Exception as e: - verbose_logger.warning( - f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}" - ) + verbose_logger.warning(f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}") return None +class _ManagedFileRow(Protocol): + unified_file_id: str + file_object: OpenAIFileObject + storage_backend: Optional[str] + storage_url: Optional[str] + created_by: Optional[str] + team_id: Optional[str] + + def model_dump(self) -> Mapping[str, object]: ... + + +class _ManagedFileTableActions(Protocol): + async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ... + + async def find_many(self, where: Mapping[str, object]) -> Sequence[_ManagedFileRow]: ... + + async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> _ManagedFileRow: ... + + async def delete(self, where: Mapping[str, str]) -> Optional[_ManagedFileRow]: ... + + +class _ManagedObjectTableActions(Protocol): + async def find_first(self, where: Mapping[str, object]) -> "Optional[PrismaManagedObjectRow]": ... + + async def find_many( + self, + where: Mapping[str, object], + take: int, + order: Union[Mapping[str, str], Sequence[Mapping[str, str]]], + cursor: Mapping[str, str] = ..., + skip: int = ..., + ) -> "Sequence[PrismaManagedObjectRow]": ... + + async def upsert( + self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] + ) -> "PrismaManagedObjectRow": ... + + +class _CursorPageArgs(TypedDict, total=False): + cursor: Mapping[str, str] + skip: int + + +def _managed_file_table(prisma_client: PrismaClient) -> _ManagedFileTableActions: + return prisma_client.db.litellm_managedfiletable + + +def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableActions: + return prisma_client.db.litellm_managedobjecttable + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes - def __init__( - self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient - ): + def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client @@ -143,9 +198,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings: Dict[str, str], user_api_key_dict: UserAPIKeyAuth, ) -> None: - verbose_logger.info( - f"Storing LiteLLM Managed File object with id={file_id} in cache" - ) + verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache") if file_object is not None: litellm_managed_file_object = LiteLLM_ManagedFileTable( unified_file_id=file_id, @@ -196,13 +249,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): f"storage_url={db_data.get('storage_url')}" ) - result = await self.prisma_client.db.litellm_managedfiletable.upsert( + result = await _managed_file_table(self.prisma_client).upsert( where={"unified_file_id": file_id}, data={"create": db_data, "update": update_data}, ) - verbose_logger.debug( - f"LiteLLM Managed File object with id={file_id} stored in db: {result}" - ) + verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") async def store_unified_object_id( self, @@ -213,9 +264,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, ) -> None: - verbose_logger.info( - f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache" - ) + verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, @@ -228,7 +277,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span=litellm_parent_otel_span, ) - await self.prisma_client.db.litellm_managedobjecttable.upsert( + await _managed_object_table(self.prisma_client).upsert( where={"unified_object_id": unified_object_id}, data={ "create": { @@ -265,9 +314,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return LiteLLM_ManagedFileTable.model_validate(result) ## CHECK DB - db_object = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + db_object = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if db_object: return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump()) @@ -277,9 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self, file_id: str, litellm_parent_otel_span: Optional[Span] = None ) -> OpenAIFileObject: ## get old value - initial_value = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + initial_value = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if initial_value is None: raise Exception(f"LiteLLM Managed File object with id={file_id} not found") ## delete old value @@ -288,15 +333,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): value=None, litellm_parent_otel_span=litellm_parent_otel_span, ) - await self.prisma_client.db.litellm_managedfiletable.delete( - where={"unified_file_id": file_id} - ) + await _managed_file_table(self.prisma_client).delete(where={"unified_file_id": file_id}) return initial_value.file_object - async def can_user_call_unified_file_id( - self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: - managed_file = await self.prisma_client.db.litellm_managedfiletable.find_first( + async def can_user_call_unified_file_id(self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + managed_file = await _managed_file_table(self.prisma_client).find_first( where={"unified_file_id": unified_file_id} ) @@ -311,13 +352,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"File not found: {unified_file_id}", ) - async def can_user_call_unified_object_id( - self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: - managed_object = ( - await self.prisma_client.db.litellm_managedobjecttable.find_first( - where={"unified_object_id": unified_object_id} - ) + async def can_user_call_unified_object_id(self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + managed_object = await _managed_object_table(self.prisma_client).find_first( + where={"unified_object_id": unified_object_id} ) if managed_object: @@ -339,34 +376,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): provider: Optional[str] = None, target_model_names: Optional[str] = None, llm_router: Optional[Router] = None, - ) -> Dict[str, Any]: + ) -> Dict[str, object]: # Provider filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the provider information # To support provider filtering, we would need to store the provider information in the encoded object ids if provider: - raise Exception( - "Filtering by 'provider' is not supported when using managed batches." - ) + raise Exception("Filtering by 'provider' is not supported when using managed batches.") # Model name filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the model name # A hash of the model name + litellm_params for the model name is encoded as the model id. This is not sufficient to reliably map the target model names to the model ids. if target_model_names: - raise Exception( - "Filtering by 'target_model_names' is not supported when using managed batches." - ) + raise Exception("Filtering by 'target_model_names' is not supported when using managed batches.") owner_filter = build_owner_filter(user_api_key_dict) if owner_filter is None: return build_list_page([]) - where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter} + where_clause: Dict[str, object] = {"file_purpose": "batch", **owner_filter} if after: - cursor_row = ( - await self.prisma_client.db.litellm_managedobjecttable.find_first( - where={**where_clause, "unified_object_id": after} - ) + cursor_row = await _managed_object_table(self.prisma_client).find_first( + where={**where_clause, "unified_object_id": after} ) if cursor_row is None: raise HTTPException( @@ -375,11 +406,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) page_size: Final = min(limit or 20, 100) - cursor_args: Dict[str, Any] = ( - {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - ) + cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( + batches = await _managed_object_table(self.prisma_client).find_many( where=where_clause, take=page_size + 1, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], @@ -389,9 +418,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): has_more = len(batches) > page_size parsed_rows: Final = tuple( - (row, batch_obj) - for row in batches[:page_size] - if (batch_obj := _parse_managed_batch_row(row)) is not None + (row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None ) unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified( raw_file_ids=frozenset( @@ -432,14 +459,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): verbose_proxy_logger=verbose_logger, user_api_key_dict=user_api_key_dict, db_batch_object=row, - unified_batch_id=_is_base64_encoded_unified_file_id( - row.unified_object_id - ), + unified_batch_id=_is_base64_encoded_unified_file_id(row.unified_object_id), ) except Exception as e: - verbose_logger.warning( - f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}" - ) + verbose_logger.warning(f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}") return None return batch_obj @@ -458,7 +481,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if owner_filter is None: return [] - file_ids = await self.prisma_client.db.litellm_managedfiletable.find_many( + file_ids = await _managed_file_table(self.prisma_client).find_many( where={ **owner_filter, "flat_model_file_ids": {"hasSome": model_object_ids}, @@ -467,27 +490,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return [ parsed_file_object.model_copy(update={"id": row.unified_file_id}) for row in file_ids - if ( - parsed_file_object := _parse_managed_file_object( - row.file_object, row.unified_file_id - ) - ) - is not None + if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None ] - async def check_managed_file_id_access( - self, data: Dict, user_api_key_dict: UserAPIKeyAuth - ) -> bool: + async def check_managed_file_id_access(self, data: Dict, user_api_key_dict: UserAPIKeyAuth) -> bool: retrieve_file_id = cast(Optional[str], data.get("file_id")) - potential_file_id = ( - _is_base64_encoded_unified_file_id(retrieve_file_id) - if retrieve_file_id - else False - ) + potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and retrieve_file_id: - if await self.can_user_call_unified_file_id( - retrieve_file_id, user_api_key_dict - ): + if await self.can_user_call_unified_file_id(retrieve_file_id, user_api_key_dict): return True else: raise HTTPException( @@ -496,9 +506,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) return False - async def check_file_ids_access( - self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth - ) -> None: + async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None: """ Check if the user has access to a list of file IDs. Only checks managed (unified) file IDs. @@ -513,9 +521,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for file_id in file_ids: is_unified_file_id = _is_base64_encoded_unified_file_id(file_id) if is_unified_file_id: - if not await self.can_user_call_unified_file_id( - file_id, user_api_key_dict - ): + if not await self.can_user_call_unified_file_id(file_id, user_api_key_dict): raise HTTPException( status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}", @@ -543,10 +549,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ### HANDLE TRANSFORMATIONS ### # Check both completion and acompletion call types - is_completion_call = ( - call_type == CallTypes.completion.value - or call_type == CallTypes.acompletion.value - ) + is_completion_call = call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value if is_completion_call: messages = data.get("messages") @@ -559,9 +562,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check if any files are stored in storage backends and need base64 conversion # This is needed for Vertex AI/Gemini which requires base64 content - is_vertex_ai = model and ( - "vertex_ai" in model or "gemini" in model.lower() - ) + is_vertex_ai = model and ("vertex_ai" in model or "gemini" in model.lower()) if is_vertex_ai: await self._convert_storage_files_to_base64( messages=messages, @@ -573,10 +574,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids, user_api_key_dict.parent_otel_span ) data["model_file_id_mapping"] = model_file_id_mapping - elif ( - call_type == CallTypes.aresponses.value - or call_type == CallTypes.responses.value - ): + elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value: # Handle managed files in responses API input and tools file_ids = [] @@ -603,23 +601,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if tools: unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools) if unified_vs_ids: - await self.check_vector_store_ids_access( - unified_vs_ids, user_api_key_dict - ) + await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict) elif call_type == CallTypes.afile_content.value: retrieve_file_id = cast(Optional[str], data.get("file_id")) - potential_file_id = ( - _is_base64_encoded_unified_file_id(retrieve_file_id) - if retrieve_file_id - else False - ) + potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and "llm_output_file_id," in potential_file_id: model_id = self.get_model_id_from_unified_file_id(potential_file_id) if model_id: data["model"] = model_id - data["file_id"] = self.get_output_file_id_from_unified_file_id( - potential_file_id - ) + data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id) elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: @@ -636,10 +626,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ): accessor_key: Optional[str] = None retrieve_object_id: Optional[str] = None - if ( - call_type == CallTypes.aretrieve_batch.value - or call_type == CallTypes.acancel_batch.value - ): + if call_type == CallTypes.aretrieve_batch.value or call_type == CallTypes.acancel_batch.value: accessor_key = "batch_id" elif ( call_type == CallTypes.acancel_fine_tuning_job.value @@ -651,32 +638,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): retrieve_object_id = cast(Optional[str], data.get(accessor_key)) potential_llm_object_id = ( - _is_base64_encoded_unified_file_id(retrieve_object_id) - if retrieve_object_id - else False + _is_base64_encoded_unified_file_id(retrieve_object_id) if retrieve_object_id else False ) if potential_llm_object_id and retrieve_object_id: ## VALIDATE USER HAS ACCESS TO THE OBJECT ## - if not await self.can_user_call_unified_object_id( - retrieve_object_id, user_api_key_dict - ): + if not await self.can_user_call_unified_object_id(retrieve_object_id, user_api_key_dict): raise HTTPException( status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the object {retrieve_object_id}", ) ## for managed batch id - get the model id - potential_model_id = get_model_id_from_unified_batch_id( - potential_llm_object_id - ) + potential_model_id = get_model_id_from_unified_batch_id(potential_llm_object_id) if potential_model_id is None: raise Exception( f"LiteLLM Managed {accessor_key} with id={retrieve_object_id} is invalid - does not contain encoded model_id." ) data["model"] = potential_model_id - data[accessor_key] = get_batch_id_from_unified_batch_id( - potential_llm_object_id - ) + data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id) elif call_type == CallTypes.acreate_fine_tuning_job.value: input_file_id = cast(Optional[str], data.get("training_file")) if input_file_id: @@ -732,24 +711,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if accessor_key: input_file_id = cast(Optional[str], kwargs.get(accessor_key)) - model_file_id_mapping = cast( - Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping") - ) + model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) # model_info may be at top-level or nested under litellm_metadata # (batch/file operations use litellm_metadata) model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) if model_id is None: model_id = cast( Optional[str], - kwargs.get("litellm_metadata", {}) - .get("model_info", {}) - .get("id", None), + kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), ) mapped_file_id: Optional[str] = None if input_file_id and model_file_id_mapping and model_id: - mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get( - model_id, None - ) + mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None) if mapped_file_id: kwargs[accessor_key] = mapped_file_id @@ -775,9 +748,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_input( - self, input: Union[str, List[Dict[str, Any]]] - ) -> List[str]: + def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]: """ Gets file ids from responses API input. @@ -809,19 +780,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): content = item.get("content") if isinstance(content, list): for content_item in content: - if ( - isinstance(content_item, dict) - and content_item.get("type") == "input_file" - ): + if isinstance(content_item, dict) and content_item.get("type") == "input_file": file_id = content_item.get("file_id") if file_id: file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_tools( - self, tools: List[Dict[str, Any]] - ) -> List[str]: + def get_file_ids_from_responses_tools(self, tools: List[Dict[str, object]]) -> List[str]: """ Gets file ids from responses API tools parameter. @@ -854,9 +820,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return file_ids - def get_vector_store_ids_from_file_search_tools( - self, tools: List[Dict[str, Any]] - ) -> List[str]: + def get_vector_store_ids_from_file_search_tools(self, tools: List[Dict[str, object]]) -> List[str]: """ Extract unified vector_store_ids from file_search tools. @@ -949,9 +913,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ), ) - async def get_model_file_id_mapping( - self, file_ids: List[str], litellm_parent_otel_span: Span - ) -> dict: + async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict: """ Get model-specific file IDs for a list of proxy file IDs. Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id @@ -981,9 +943,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Get all cache keys matching the pattern file_id:* for file_id in litellm_managed_file_ids: # Search for any cache key starting with this file_id - unified_file_object = await self.get_unified_file_id( - file_id, litellm_parent_otel_span - ) + unified_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) if unified_file_object: file_id_mapping[file_id] = unified_file_object.model_mappings @@ -1001,9 +961,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): raise Exception("LLM Router not initialized. Ensure models added to proxy.") responses = [] for model in target_model_names_list: - individual_response = await llm_router.acreate_file( - model=model, **_create_file_request - ) + individual_response = await llm_router.acreate_file(model=model, **_create_file_request) responses.append(individual_response) return responses @@ -1034,9 +992,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings: Dict[str, str] = {} for file_object in responses: - model_file_id_mapping = file_object._hidden_params.get( - "model_file_id_mapping" - ) + model_file_id_mapping = file_object._hidden_params.get("model_file_id_mapping") if model_file_id_mapping and isinstance(model_file_id_mapping, dict): model_mappings.update(model_file_id_mapping) @@ -1051,17 +1007,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Emit Prometheus metrics for managed file creation prom_logger = self._get_prometheus_logger() if prom_logger: - first_model = ( - target_model_names_list[0] if target_model_names_list else None - ) + first_model = target_model_names_list[0] if target_model_names_list else None first_provider = "" if responses: - first_provider = ( - getattr(responses[0], "_hidden_params", {}).get( - "custom_llm_provider" - ) - or "" - ) + first_provider = getattr(responses[0], "_hidden_params", {}).get("custom_llm_provider") or "" prom_logger.record_managed_file_created( model=first_model or "", api_provider=first_provider, @@ -1104,9 +1053,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) # Convert to URL-safe base64 and strip padding - base64_unified_file_id = ( - base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") - ) + base64_unified_file_id = base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") ## CREATE RESPONSE OBJECT @@ -1123,46 +1070,26 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response - def get_unified_generic_response_id( - self, model_id: str, generic_response_id: str - ) -> str: - unified_generic_response_id = ( - SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format( - model_id, generic_response_id - ) - ) - return ( - base64.urlsafe_b64encode(unified_generic_response_id.encode()) - .decode() - .rstrip("=") + def get_unified_generic_response_id(self, model_id: str, generic_response_id: str) -> str: + unified_generic_response_id = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format( + model_id, generic_response_id ) + return base64.urlsafe_b64encode(unified_generic_response_id.encode()).decode().rstrip("=") def get_unified_batch_id(self, batch_id: str, model_id: str) -> str: - unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( - model_id, batch_id - ) + unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") - def get_unified_output_file_id( - self, output_file_id: str, model_id: str, model_name: Optional[str] - ) -> str: - deterministic_uuid: Final = uuid5( - uuid5(NAMESPACE_URL, model_id), output_file_id - ) - unified_output_file_id = ( - SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - "application/json", - str(deterministic_uuid), - model_name or "", - output_file_id, - model_id, - ) - ) - return ( - base64.urlsafe_b64encode(unified_output_file_id.encode()) - .decode() - .rstrip("=") + def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str: + deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id) + unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", + str(deterministic_uuid), + model_name or "", + output_file_id, + model_id, ) + return base64.urlsafe_b64encode(unified_output_file_id.encode()).decode().rstrip("=") def get_model_id_from_unified_file_id(self, file_id: str) -> str: return file_id.split("llm_output_file_model_id,")[1].split(";")[0] @@ -1170,59 +1097,39 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): def get_output_file_id_from_unified_file_id(self, file_id: str) -> str: marker = "llm_output_file_id," if marker not in file_id: - raise ValueError( - f"Unified id does not contain {marker!r}: {file_id[:80]!r}" - ) + raise ValueError(f"Unified id does not contain {marker!r}: {file_id[:80]!r}") return file_id.split(marker, 1)[1].split(";")[0] async def async_post_call_success_hook( self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes - ) -> Any: + ) -> LLMResponseTypes: if isinstance(response, LiteLLMBatch): ## Check if unified_file_id is in the response - unified_file_id = response._hidden_params.get( - "unified_file_id" - ) # managed file id - unified_batch_id = response._hidden_params.get( - "unified_batch_id" - ) # managed batch id + unified_file_id = response._hidden_params.get("unified_file_id") # managed file id + unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id model_id = cast(Optional[str], response._hidden_params.get("model_id")) model_name = cast(Optional[str], response._hidden_params.get("model_name")) resolved_model_name = resolve_managed_output_file_model_name( - unified_input_file_id=unified_file_id - if isinstance(unified_file_id, str) - else response.input_file_id, + unified_input_file_id=unified_file_id if isinstance(unified_file_id, str) else response.input_file_id, fallback_model_name=model_name, ) original_response_id = response.id if (unified_batch_id or unified_file_id) and model_id: - response.id = self.get_unified_batch_id( - batch_id=response.id, model_id=model_id - ) + response.id = self.get_unified_batch_id(batch_id=response.id, model_id=model_id) # Handle both output_file_id and error_file_id for file_attr in ["output_file_id", "error_file_id"]: file_id_value = getattr(response, file_attr, None) if file_id_value and model_id: - decoded_output_file_id = _is_base64_encoded_unified_file_id( - file_id_value - ) - if ( - decoded_output_file_id - and "llm_output_file_id," in decoded_output_file_id - ): - provider_file_id = ( - self.get_output_file_id_from_unified_file_id( - decoded_output_file_id - ) - ) + decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value) + if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id: + provider_file_id = self.get_output_file_id_from_unified_file_id(decoded_output_file_id) unified_file_id = file_id_value elif decoded_output_file_id: verbose_logger.warning( - f"Skipping {file_attr}={file_id_value!r}: " - "unified id is not a managed file output id" + f"Skipping {file_attr}={file_id_value!r}: unified id is not a managed file output id" ) continue else: @@ -1241,23 +1148,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Import module and use getattr for better testability with mocks import litellm.proxy.proxy_server as proxy_server_module - _llm_router = getattr( - proxy_server_module, "llm_router", None - ) + _llm_router = getattr(proxy_server_module, "llm_router", None) if _llm_router is not None and model_id: - _creds = ( - _llm_router.get_deployment_credentials_with_provider( - model_id - ) - or {} - ) + _creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {} file_object = await litellm.afile_retrieve( file_id=provider_file_id, **_creds, ) else: file_object = await litellm.afile_retrieve( - custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", # type: ignore[arg-type] + custom_llm_provider=model_name.split("/")[0] + if model_name and "/" in model_name + else "openai", # type: ignore[arg-type] file_id=provider_file_id, ) verbose_logger.debug( @@ -1311,9 +1213,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): elif isinstance(response, LiteLLMFineTuningJob): ## Check if unified_file_id is in the response - unified_file_id = response._hidden_params.get( - "unified_file_id" - ) # managed file id + unified_file_id = response._hidden_params.get("unified_file_id") # managed file id unified_finetuning_job_id = response._hidden_params.get( "unified_finetuning_job_id" ) # managed finetuning job id @@ -1321,9 +1221,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_name = cast(Optional[str], response._hidden_params.get("model_name")) original_response_id = response.id if (unified_file_id or unified_finetuning_job_id) and model_id: - response.id = self.get_unified_generic_response_id( - model_id=model_id, generic_response_id=response.id - ) + response.id = self.get_unified_generic_response_id(model_id=model_id, generic_response_id=response.id) await self.store_unified_object_id( unified_object_id=response.id, file_object=response, @@ -1338,9 +1236,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ ## check if file object if hasattr(response, "data") and isinstance(response.data, list): - if all( - isinstance(file_object, FileObject) for file_object in response.data - ): + if all(isinstance(file_object, FileObject) for file_object in response.data): ## Get all file id's ## Check which file id's were created by the user ## Filter the response to only include the files created by the user @@ -1349,9 +1245,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_object.id for file_object in cast(List[FileObject], response.data) # type: ignore ] - user_created_file_ids = await self.get_user_created_file_ids( - user_api_key_dict, file_ids - ) + user_created_file_ids = await self.get_user_created_file_ids(user_api_key_dict, file_ids) ## Filter the response to only include the files created by the user response.data = user_created_file_ids # type: ignore return response @@ -1359,11 +1253,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response async def afile_retrieve( - self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router=None + self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router: Optional[Router] = None ) -> OpenAIFileObject: - stored_file_object = await self.get_unified_file_id( - file_id, litellm_parent_otel_span - ) + stored_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) # Case 1 : This is not a managed file if not stored_file_object: @@ -1386,21 +1278,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) try: - model_id, model_file_id = next( - iter(stored_file_object.model_mappings.items()) - ) - credentials = ( - llm_router.get_deployment_credentials_with_provider(model_id) or {} - ) - response = await litellm.afile_retrieve( - file_id=model_file_id, **credentials - ) + model_id, model_file_id = next(iter(stored_file_object.model_mappings.items())) + credentials = llm_router.get_deployment_credentials_with_provider(model_id) or {} + response = await litellm.afile_retrieve(file_id=model_file_id, **credentials) response.id = file_id # Replace with unified ID return response except Exception as e: - raise Exception( - f"Failed to retrieve file {file_id} from provider: {str(e)}" - ) from e + raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e async def afile_list( self, @@ -1437,12 +1321,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return False except Exception as e: - verbose_logger.warning( - f"Error checking batch polling configuration: {e}. Assuming disabled." - ) + verbose_logger.warning(f"Error checking batch polling configuration: {e}. Assuming disabled.") return False - async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, Any]]: + async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, object]]: """ Find batches that reference this file and still need cost tracking. Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost. @@ -1458,9 +1340,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Get model-specific file IDs for this unified file ID if it's a managed file try: - model_file_id_mapping = await self.get_model_file_id_mapping( - [file_id], litellm_parent_otel_span=None - ) + model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span=None) if model_file_id_mapping and file_id in model_file_id_mapping: # Add all provider file IDs for this unified file @@ -1468,8 +1348,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids_to_check.extend(provider_file_ids) except Exception as e: verbose_logger.debug( - f"Could not get model file ID mapping for {file_id}: {e}. " - f"Will only check unified file ID." + f"Could not get model file ID mapping for {file_id}: {e}. Will only check unified file ID." ) MAX_MATCHES_TO_RETURN = 10 @@ -1487,11 +1366,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for batch in batches: try: # Parse the batch file_object to check for file references - batch_data = ( - json.loads(batch.file_object) - if isinstance(batch.file_object, str) - else batch.file_object - ) + batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object # Extract file IDs from batch # Batches typically reference the unified file ID in input_file_id @@ -1500,9 +1375,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): output_file_id = batch_data.get("output_file_id") error_file_id = batch_data.get("error_file_id") - referenced_file_ids = [ - fid for fid in [input_file_id, output_file_id, error_file_id] if fid - ] + referenced_file_ids = [fid for fid in [input_file_id, output_file_id, error_file_id] if fid] # Check if any referenced file ID matches the file we're trying to delete if any(ref_id in file_ids_to_check for ref_id in referenced_file_ids): @@ -1514,9 +1387,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ) except Exception as e: - verbose_logger.warning( - f"Error parsing batch object {batch.unified_object_id}: {e}" - ) + verbose_logger.warning(f"Error parsing batch object {batch.unified_object_id}: {e}") continue return referencing_batches @@ -1545,21 +1416,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if referencing_batches: # File is referenced by non-terminal batches and polling is enabled - MAX_BATCHES_IN_ERROR = ( - 5 # Limit batches shown in error message for readability - ) + MAX_BATCHES_IN_ERROR = 5 # Limit batches shown in error message for readability # Show up to MAX_BATCHES_IN_ERROR in the error message batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR] - batch_statuses = [ - f"{b['batch_id']}: {b['status']}" for b in batches_to_show - ] + batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show] # Determine the count message count_message = f"{len(referencing_batches)}" - if ( - len(referencing_batches) >= 10 - ): # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file + if len(referencing_batches) >= 10: # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file count_message = "10+" error_message = ( @@ -1600,23 +1465,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): await self._check_file_deletion_allowed(file_id) # file_id = convert_b64_uid_to_unified_uid(file_id) - model_file_id_mapping = await self.get_model_file_id_mapping( - [file_id], litellm_parent_otel_span - ) + model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) delete_response = None specific_model_file_id_mapping = model_file_id_mapping.get(file_id) if specific_model_file_id_mapping: # Remove conflicting keys from data to avoid duplicate keyword arguments - filtered_data = { - k: v for k, v in data.items() if k not in ("model", "file_id") - } + filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")} for model_id, model_file_id in specific_model_file_id_mapping.items(): delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore - stored_file_object = await self.delete_unified_file_id( - file_id, litellm_parent_otel_span - ) + stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span) # Record successful deletion metric only on actual success if stored_file_object or delete_response: @@ -1643,9 +1502,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): Get the content of a file from first model that has it """ model_file_id_mapping = data.pop("model_file_id_mapping", None) - model_file_id_mapping = ( - model_file_id_mapping - or await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) + model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping( + [file_id], litellm_parent_otel_span ) specific_model_file_id_mapping = model_file_id_mapping.get(file_id) @@ -1658,13 +1516,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # against the deployment's configured bucket, which they only # trust from this immutable server-side snapshot, never from # request params. - credentials = llm_router.get_deployment_credentials_with_provider( - model_id=model_id - ) + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) if credentials is not None: - data["_litellm_internal_model_credentials"] = cast( - Dict, MappingProxyType(dict(credentials)) - ) + data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials))) else: data.pop("_litellm_internal_model_credentials", None) return await llm_router.afile_content(model=model_id, file_id=provider_file_id, **data) # type: ignore @@ -1699,9 +1553,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check database for storage backend info # IMPORTANT: The database stores the base64 encoded unified_file_id (not the decoded version) # So we query with the original file_id (which is base64 encoded) - db_file = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + db_file = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if not db_file or not db_file.storage_backend or not db_file.storage_url: continue @@ -1727,22 +1579,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_content = await storage_backend.download_file(storage_url) # Determine content type from file object - content_type = self._get_content_type_from_file_object( - db_file.file_object - ) + content_type = self._get_content_type_from_file_object(db_file.file_object) # Convert to base64 base64_data = base64.b64encode(file_content).decode("utf-8") base64_data_uri = f"data:{content_type};base64,{base64_data}" # Update messages to use base64 instead of file_id - self._update_messages_with_base64_data( - messages, file_id, base64_data_uri, content_type - ) + self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type) except Exception as e: - verbose_logger.exception( - f"Error converting file {file_id} from storage backend to base64: {str(e)}" - ) + verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {str(e)}") # Continue with other files even if one fails continue diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 339da998d56..024e8c179c2 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,16 +6,33 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from collections.abc import AsyncIterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Any, Final, Protocol, cast, runtime_checkable from uuid import uuid4 +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) +_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_TEXT_ADAPTER: Final = TypeAdapter(str) + + +@runtime_checkable +class _SupportsModelDump(Protocol): + def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ... + + +@runtime_checkable +class _SupportsPydanticDict(Protocol): + def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ... + class PydanticAITransformation: """ @@ -28,7 +45,7 @@ class PydanticAITransformation: """ @staticmethod - def _remove_none_values(obj: Any) -> Any: + def _remove_none_values(obj: object) -> object: """ Recursively remove None values from a dict/list structure. @@ -42,14 +59,18 @@ class PydanticAITransformation: Cleaned object with None values removed """ if isinstance(obj, dict): - return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None} + typed_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(obj) + return {k: PydanticAITransformation._remove_none_values(v) for k, v in typed_dict.items() if v is not None} elif isinstance(obj, list): - return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None] + typed_list: Final = _LIST_ADAPTER.validate_python(obj) + return [PydanticAITransformation._remove_none_values(item) for item in typed_list if item is not None] else: return obj @staticmethod - def _params_to_dict(params: Any) -> dict[str, Any]: + def _params_to_dict( + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", + ) -> Mapping[str, object]: """ Convert params to a dict, handling Pydantic models. @@ -59,10 +80,10 @@ class PydanticAITransformation: Returns: Dict representation of params """ - if hasattr(params, "model_dump"): + if isinstance(params, _SupportsModelDump): # Pydantic v2 model return params.model_dump(mode="python", exclude_none=True) - elif hasattr(params, "dict"): + elif isinstance(params, _SupportsPydanticDict): # Pydantic v1 model return params.dict(exclude_none=True) elif isinstance(params, dict): @@ -75,12 +96,12 @@ class PydanticAITransformation: async def _poll_for_completion( client: AsyncHTTPHandler, endpoint: str, - task_id: str, + task_id: object, request_id: str, max_attempts: int = 30, poll_interval: float = 0.5, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Poll for task completion using tasks/get method. @@ -112,10 +133,10 @@ class PydanticAITransformation: }, ) response.raise_for_status() - poll_data = response.json() + poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) - result = poll_data.get("result", {}) - status = result.get("status", {}) + result = _STR_KEY_DICT_ADAPTER.validate_python(poll_data.get("result", {})) + status = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {})) state = status.get("state", "") verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state) @@ -133,10 +154,10 @@ class PydanticAITransformation: async def _send_and_poll_raw( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -153,14 +174,16 @@ class PydanticAITransformation: Raw Pydantic AI task response (with history/artifacts) """ # Convert params to dict if it's a Pydantic model - params_dict = PydanticAITransformation._params_to_dict(params) - # Remove None values - FastA2A doesn't accept null for optional fields - params_dict = PydanticAITransformation._remove_none_values(params_dict) + params_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python( + PydanticAITransformation._remove_none_values(PydanticAITransformation._params_to_dict(params)) + ) # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI if "message" in params_dict: - params_dict["message"]["kind"] = "message" + message_value: Final = _ANY_KEY_DICT_ADAPTER.validate_python(params_dict["message"]) + message_value["kind"] = "message" + params_dict["message"] = message_value # Build A2A JSON-RPC request using message/send method for FastA2A compatibility a2a_request: Final = { @@ -189,11 +212,11 @@ class PydanticAITransformation: }, ) response.raise_for_status() - response_data = response.json() + response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) # Check if task is already completed - result: Final = response_data.get("result", {}) - status: Final = result.get("status", {}) + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) + status: Final = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {})) state: Final = status.get("state", "") if state != "completed": @@ -217,10 +240,10 @@ class PydanticAITransformation: async def send_non_streaming_request( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a non-streaming A2A request to Pydantic AI agent and wait for completion. @@ -253,10 +276,10 @@ class PydanticAITransformation: async def send_and_get_raw_response( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -282,9 +305,9 @@ class PydanticAITransformation: @staticmethod def _transform_to_a2a_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform Pydantic AI task response to standard A2A non-streaming format. @@ -328,7 +351,7 @@ class PydanticAITransformation: } @staticmethod - def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]: + def _extract_response_text(response_data: Mapping[str, object]) -> tuple[object, object, Sequence[object]]: """ Extract response text from completed task response. @@ -342,52 +365,53 @@ class PydanticAITransformation: Returns: Tuple of (full_text, message_id, parts) """ - result: Final = response_data.get("result", {}) + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) # Try to extract from artifacts first (preferred for results) artifacts: Final = result.get("artifacts", []) if artifacts: - for artifact in artifacts: - parts = artifact.get("parts", []) + for artifact in _LIST_ADAPTER.validate_python(artifacts): + parts = _LIST_ADAPTER.validate_python(_STR_KEY_DICT_ADAPTER.validate_python(artifact).get("parts", [])) for part in parts: - if part.get("kind") == "text": - text = part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + text = part_dict.get("text", "") if text: return text, str(uuid4()), parts # Fall back to history - get the last agent message - history: Final = result.get("history", []) + history: Final = _LIST_ADAPTER.validate_python(result.get("history", [])) for msg in reversed(history): - if msg.get("role") == "agent": - parts = msg.get("parts", []) - message_id = msg.get("messageId", str(uuid4())) + if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "agent": + parts = _LIST_ADAPTER.validate_python(msg_dict.get("parts", [])) + message_id = msg_dict.get("messageId", str(uuid4())) full_text = "" for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", "")) if full_text: return full_text, message_id, parts # Fall back to message field (original format) message: Final = result.get("message", {}) if message: - parts = message.get("parts", []) - message_id = message.get("messageId", str(uuid4())) + message_dict: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) + parts = _LIST_ADAPTER.validate_python(message_dict.get("parts", [])) + message_id = message_dict.get("messageId", str(uuid4())) full_text = "" for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", "")) return full_text, message_id, parts return "", str(uuid4()), [] @staticmethod async def fake_streaming_from_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """ Convert a non-streaming A2A response into fake streaming chunks. @@ -410,12 +434,12 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Extract input message from raw response for history - result: Final = response_data.get("result", {}) - history: Final = result.get("history", []) - input_message = {} + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) + history: Final = _LIST_ADAPTER.validate_python(result.get("history", [])) + input_message = _STR_KEY_DICT_ADAPTER.validate_python({}) for msg in history: - if msg.get("role") == "user": - input_message = msg + if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "user": + input_message = msg_dict break # Generate IDs for streaming events @@ -426,45 +450,49 @@ class PydanticAITransformation: # 1. Emit initial task event (kind: "task", status: "submitted") # Format matches A2ACompletionBridgeTransformation.create_task_event - task_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "history": [ - { - "contextId": context_id, - "kind": "message", - "messageId": input_message_id, - "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), - "role": "user", - "taskId": task_id, - } - ], - "id": task_id, - "kind": "task", - "status": { - "state": "submitted", + task_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "history": [ + { + "contextId": context_id, + "kind": "message", + "messageId": input_message_id, + "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), + "role": "user", + "taskId": task_id, + } + ], + "id": task_id, + "kind": "task", + "status": { + "state": "submitted", + }, }, - }, - } + } + ) yield task_event # 2. Emit status update (kind: "status-update", status: "working") # Format matches A2ACompletionBridgeTransformation.create_status_update_event - working_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "final": False, - "kind": "status-update", - "status": { - "state": "working", + working_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": False, + "kind": "status-update", + "status": { + "state": "working", + }, + "taskId": task_id, }, - "taskId": task_id, - }, - } + } + ) yield working_event # Small delay to simulate processing @@ -473,29 +501,32 @@ class PydanticAITransformation: # 3. Emit artifact update chunks (kind: "artifact-update") # Format matches A2ACompletionBridgeTransformation.create_artifact_update_event if full_text: + full_text_str: Final = _TEXT_ADAPTER.validate_python(full_text) # Split text into chunks - for i in range(0, len(full_text), chunk_size): - chunk_text = full_text[i : i + chunk_size] - is_last_chunk = (i + chunk_size) >= len(full_text) + for i in range(0, len(full_text_str), chunk_size): + chunk_text = full_text_str[i : i + chunk_size] + is_last_chunk = (i + chunk_size) >= len(full_text_str) - artifact_event = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "kind": "artifact-update", - "taskId": task_id, - "artifact": { - "artifactId": artifact_id, - "parts": [ - { - "kind": "text", - "text": chunk_text, - } - ], + artifact_event = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "kind": "artifact-update", + "taskId": task_id, + "artifact": { + "artifactId": artifact_id, + "parts": [ + { + "kind": "text", + "text": chunk_text, + } + ], + }, }, - }, - } + } + ) yield artifact_event # Add delay between chunks (except for last chunk) @@ -503,19 +534,21 @@ class PydanticAITransformation: await asyncio.sleep(delay_ms / 1000.0) # 4. Emit final status update (kind: "status-update", status: "completed", final: true) - completed_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "final": True, - "kind": "status-update", - "status": { - "state": "completed", + completed_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": True, + "kind": "status-update", + "status": { + "state": "completed", + }, + "taskId": task_id, }, - "taskId": task_id, - }, - } + } + ) yield completed_event verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 9748db2dcd2..7abbf0c96e5 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -42,7 +42,7 @@ from litellm.types.integrations.websearch_interception import ( WebSearchInterceptionConfig, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.utils import AgenticLoopParams, CallTypes, LlmProviders from litellm.utils import ProviderConfigManager if TYPE_CHECKING: @@ -265,7 +265,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None # Check if request has tools with native web_search - tools: Final = kwargs.get("tools") + tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools") if not tools: return None @@ -314,7 +314,9 @@ class WebSearchInterceptionLogger(CustomLogger): return kwargs - def _convert_responses_tools(self, kwargs: Mapping[str, object], tools: list[dict[str, object]]) -> dict | None: + def _convert_responses_tools( + self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]] + ) -> dict[str, object] | None: """Convert Responses API web search tools to the LiteLLM standard function tool.""" if not any(is_web_search_tool_responses(tool) for tool in tools): return None @@ -379,7 +381,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _tool_name(tool: dict[str, Any]) -> str | None: + def _tool_name(tool: Mapping[str, object]) -> object: """Effective tool name, handling OpenAI ``function`` wrapper shape.""" fn: Final = tool.get("function") if tool.get("type") == "function" and isinstance(fn, dict): @@ -1271,7 +1273,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: - agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {}) + agentic_params: Final[AgenticLoopParams] = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = agentic_params.get("model", model) verbose_logger.debug( "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]", diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 858d10df53b..d68bdc4a250 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,5 +1,6 @@ import asyncio import json +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm @@ -20,11 +21,24 @@ from .litellm_logging import Logging as LiteLLMLogging if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection + from litellm.types.guardrails import GuardrailEventHooks + CLIENT_CONNECTION_CLASS = ClientConnection else: CLIENT_CONNECTION_CLASS = Any +class _ClientWebSocketExceptions(Protocol): + ConnectionClosed: type[Exception] + + +class _ClientWebSocket(Protocol): + exceptions: _ClientWebSocketExceptions + + async def send_text(self, data: str) -> None: ... + async def receive_text(self) -> str: ... + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... def normalize(self, event: dict) -> dict: ... @@ -48,13 +62,13 @@ class RealTimeStreaming: logging_obj: LiteLLMLogging, provider_config: BaseRealtimeConfig | None = None, model: str = "", - user_api_key_dict: Any | None = None, + user_api_key_dict: object | None = None, request_data: dict | None = None, backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, ): - self.websocket = websocket + self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.messages: list[OpenAIRealtimeEvents] = [] @@ -127,7 +141,7 @@ class RealTimeStreaming: ] ) _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"]) - _AUDIO_FORMAT_MAP: dict[str, dict[str, Any]] = { + _AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = { "pcm16": {"type": "audio/pcm", "rate": 24000}, "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000}, "g711_alaw": {"type": "audio/G711-alaw", "rate": 8000}, @@ -281,6 +295,7 @@ class RealTimeStreaming: if event_obj.get("type") != "response.done": return response: Final = cast(dict[str, Any], event_obj.get("response", {})) + item: Mapping[str, object] for item in response.get("output", []): if item.get("type") == "function_call": self.tool_calls.append( @@ -384,7 +399,7 @@ class RealTimeStreaming: return message try: - message_obj: Final = json.loads(message) + message_obj: Final[Mapping[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return message @@ -487,7 +502,7 @@ class RealTimeStreaming: if self._backend_setup_complete and not self._flushing_pending_messages_until_setup: return False try: - msg_obj: Final = json.loads(message) + msg_obj: Final[Mapping[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return False return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES @@ -555,7 +570,7 @@ class RealTimeStreaming: def _event_to_client_json(self, event: dict) -> str: return json.dumps(self._normalize_event_for_ga_client(event)) - async def _send_event_to_client(self, event: Any, event_str: str) -> bool: + async def _send_event_to_client(self, event: object, event_str: str) -> bool: if self._should_drop_event_from_client(event): return False if isinstance(event, dict): @@ -595,12 +610,12 @@ class RealTimeStreaming: def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" - turn_detection: Final[dict[str, Any]] = { + turn_detection: Final[dict[str, str | bool]] = { "type": "server_vad", "create_response": False, } if self._backend_uses_beta_protocol: - session: dict[str, Any] = {"turn_detection": turn_detection} + session: dict[str, object] = {"turn_detection": turn_detection} else: session = { "type": "realtime", @@ -654,7 +669,7 @@ class RealTimeStreaming: def _has_realtime_guardrails_for_event_hooks( self, - event_hooks: list[Any], + event_hooks: Sequence["GuardrailEventHooks"], ) -> bool: """Return True if any callback would run for one of ``event_hooks``.""" from litellm.integrations.custom_guardrail import CustomGuardrail @@ -699,7 +714,7 @@ class RealTimeStreaming: transcript: str, item_id: str | None = None, pre_block_backend_message: str | None = None, - event_hooks: list[Any] | None = None, + event_hooks: Sequence["GuardrailEventHooks"] | None = None, ) -> bool: """ Run registered guardrails on realtime text (transcript, user message, tool output). @@ -753,7 +768,7 @@ class RealTimeStreaming: raise # Extract the human-readable error from the detail dict (HTTPException) # or fall back to str(e) for plain ValueError. - detail = getattr(e, "detail", None) + detail: object | None = getattr(e, "detail", None) if isinstance(detail, dict): safe_msg = detail.get("error") or str(e) elif detail is not None: @@ -826,7 +841,7 @@ class RealTimeStreaming: return True return False - async def _handle_provider_config_message(self, raw_response) -> None: + async def _handle_provider_config_message(self, raw_response: str) -> None: """Process a backend message when a provider_config is set (transformed path).""" returned_object: Final = self.provider_config.transform_realtime_response( raw_response, @@ -910,7 +925,7 @@ class RealTimeStreaming: await self._send_event_to_client(event, event_str) @staticmethod - def _parse_backend_event(raw_response: str) -> dict | None: + def _parse_backend_event(raw_response: str) -> dict[str, object] | None: """Parse a backend frame once. Returns None for non-JSON or non-object frames.""" try: event: Final = json.loads(raw_response) @@ -1020,7 +1035,7 @@ class RealTimeStreaming: objects and any test doubles that expose a .scope dict. """ try: - headers: Final = websocket.scope.get("headers", []) + headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", []) for name, value in headers: if isinstance(name, bytes): name = name.decode("latin-1") @@ -1071,9 +1086,9 @@ class RealTimeStreaming: session["output_modalities"] = ["text"] # 3-7. Lift flat audio fields into the nested audio object - audio: Final[dict[str, Any]] = {} - inp: Final[dict[str, Any]] = {} - out: Final[dict[str, Any]] = {} + audio: Final[dict[str, object]] = {} + inp: Final[dict[str, object]] = {} + out: Final[dict[str, object]] = {} # voice → audio.output.voice if "voice" in session: @@ -1190,7 +1205,7 @@ class RealTimeStreaming: # model; check them with the same guardrail used for # user text so an attacker cannot smuggle blocked # content into a function_call_output. - output = item.get("output", "") + output: object = item.get("output", "") output_text = output if isinstance(output, str) else json.dumps(output) if output_text: # Build the sanitized function_call_output up @@ -1241,7 +1256,7 @@ class RealTimeStreaming: # interaction turn. continue elif item.get("role") == "user": - content_list = item.get("content", []) + content_list: Sequence[object] = item.get("content", []) texts = [ c.get("text", "") for c in content_list @@ -1280,7 +1295,7 @@ class RealTimeStreaming: and not self._guardrail_turn_detection_update_sent and self._has_audio_transcription_guardrails() ): - session = msg_obj.setdefault("session", {}) + session: object = msg_obj.setdefault("session", {}) if isinstance(session, dict): existing_td = session.get("turn_detection") if not isinstance(existing_td, dict): diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 3f967e29002..886ba6a3a18 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -3,7 +3,7 @@ import time from collections.abc import Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, Union, cast from litellm._logging import verbose_logger from litellm.types.llms.openai import ( @@ -30,6 +30,7 @@ from litellm.types.utils import ( from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, ) @@ -39,6 +40,60 @@ if TYPE_CHECKING: ) +class _ThinkingBlockFragment(TypedDict, total=False): + type: str | None + data: str | None + thinking: str | None + signature: str | None + + +class _ThinkingDelta(TypedDict, total=False): + thinking_blocks: Sequence[_ThinkingBlockFragment] + + +class _ThinkingChoice(TypedDict, total=False): + delta: _ThinkingDelta + + +class _ThinkingChunk(TypedDict): + choices: Sequence[_ThinkingChoice] + + +class _ContentChoice(TypedDict, total=False): + delta: Mapping[str, str | None] + + +class _ContentChunk(TypedDict): + choices: Sequence[_ContentChoice] + + +class _AudioDelta(TypedDict, total=False): + audio: ChatCompletionAudioDelta | None + + +class _AudioChoice(TypedDict, total=False): + delta: _AudioDelta + + +class _AudioChunk(TypedDict): + choices: Sequence[_AudioChoice] + + +class _UsageBearingChunk(TypedDict, total=False): + usage: Usage | None + _hidden_params: Mapping[str, str] + + +class _UsageSummary(TypedDict): + prompt_tokens: int | None + completion_tokens: int | None + cache_creation_input_tokens: int | None + cache_read_input_tokens: int | None + completion_tokens_details: CompletionTokensDetails | None + prompt_tokens_details: PromptTokensDetailsWrapper | None + cost: float | None + + def capture_cache_creation_token_details( prompt_tokens_details: PromptTokensDetailsWrapper | None, current: CacheCreationTokenDetails | None, @@ -78,7 +133,7 @@ class ChunkProcessor: return [] first_chunk: Final = chunks[0] - first_hidden_params: dict[str, Any] = {} + first_hidden_params: dict[str, object] = {} if isinstance(first_chunk, dict): candidate = first_chunk.get("_hidden_params", {}) if isinstance(candidate, dict): @@ -115,8 +170,8 @@ class ChunkProcessor: @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, - chunks: list[Any], - logging_obj: Any | None = None, + chunks: list[object], + logging_obj: "Logging | None" = None, ) -> None: if not chunks: return @@ -456,7 +511,7 @@ class ChunkProcessor: ) def get_combined_content( - self, chunks: list[dict[str, Any]], delta_key: str = "content" + self, chunks: Sequence["_ContentChunk"], delta_key: str = "content" ) -> ChatCompletionAssistantContentValue: content_list: Final[list[str]] = [] for chunk in chunks: @@ -475,7 +530,7 @@ class ChunkProcessor: return combined_content def get_combined_thinking_content( - self, chunks: list[dict[str, Any]] + self, chunks: Sequence["_ThinkingChunk"] ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, @@ -532,10 +587,10 @@ class ChunkProcessor: return thinking_blocks return None - def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue: + def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") - def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse: + def get_combined_audio_content(self, chunks: Sequence["_AudioChunk"]) -> ChatCompletionAudioResponse: base64_data_list: Final[list[str]] = [] transcript_list: Final[list[str]] = [] expires_at: int | None = None @@ -544,7 +599,7 @@ class ChunkProcessor: for chunk in chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta") or {} + delta: _AudioDelta = choice.get("delta") or {} audio: ChatCompletionAudioDelta | None = delta.get("audio") if audio is not None: for k, v in audio.items(): @@ -565,7 +620,7 @@ class ChunkProcessor: id=id, ) - def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict: + def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> "_UsageSummary": prompt_tokens = 0 completion_tokens = 0 ## anthropic prompt caching information ## @@ -623,8 +678,8 @@ class ChunkProcessor: return reasoning_tokens @staticmethod - def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None: - usage_chunk: Usage | dict[str, Any] | None = None + def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None: + usage_chunk: Usage | None = None if hasattr(chunk, "usage") and chunk.usage is not None: usage_chunk = chunk.usage elif "usage" in chunk: @@ -640,7 +695,7 @@ class ChunkProcessor: def _calculate_usage_per_chunk( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], ) -> "UsagePerChunk": from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -721,13 +776,7 @@ class ChunkProcessor: "web_search_requests", ) - prompt_tokens_details = ( - cast( - PromptTokensDetailsWrapper | None, - usage_chunk_dict["prompt_tokens_details"], - ) - or prompt_tokens_details - ) + prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details cache_creation_token_details = capture_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details @@ -758,7 +807,7 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], completion_tokens: int, completion_usage_updates: int, ) -> int: @@ -797,7 +846,7 @@ class ChunkProcessor: def calculate_usage( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, messages: list | None = None, @@ -851,8 +900,8 @@ class ChunkProcessor: setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens) # for anthropic if completion_tokens_details is not None: if isinstance(completion_tokens_details, CompletionTokensDetails): - returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - **completion_tokens_details.model_dump() + returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper.model_validate( + completion_tokens_details.model_dump() ) else: returned_usage.completion_tokens_details = completion_tokens_details diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a9751489473..36f3e875a7e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,8 +1,9 @@ -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from typing import ( TYPE_CHECKING, Any, Final, + TypeAlias, cast, ) @@ -33,8 +34,12 @@ if TYPE_CHECKING: # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge. ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config"}) +_AnthropicMessages: TypeAlias = "list[dict[str, object]]" +_AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None" +_ContextManagementSpec: TypeAlias = "dict[str, object] | list[dict[str, object]] | None" -def _messages_have_compaction_block(messages: list[dict]) -> bool: + +def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: content = msg.get("content") @@ -54,8 +59,10 @@ def _proxy_router_fallback() -> "Router | None": return _proxy_router -def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: - """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise. +def _extract_proxy_litellm_metadata( + kwargs: Mapping[str, object], +) -> "tuple[dict[str, object], UserAPIKeyAuth | None] | tuple[None, None]": + """Return ``(kwargs["litellm_metadata"], its user_api_key_auth)`` when it's a dict; ``(None, None)`` otherwise. The proxy attaches its auth/spend-attribution fields (``user_api_key``, ``user_api_key_team_id``, ``litellm_call_id``, the full ``UserAPIKeyAuth`` @@ -68,18 +75,19 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | """ litellm_metadata: Final = kwargs.get("litellm_metadata") if not isinstance(litellm_metadata, dict): - return None - return litellm_metadata + return None, None + user_api_key_auth: Final[UserAPIKeyAuth | None] = litellm_metadata.get("user_api_key_auth") + return litellm_metadata, user_api_key_auth async def _prepare_context_managed_request( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, + messages: _AnthropicMessages, + tools: list[dict[str, object]] | None, + system: _AnthropicSystem, + context_management_spec: _ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, llm_router: "Router | None", user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -102,11 +110,11 @@ async def _prepare_context_managed_request( if polyfill_will_run: history_result: PolyfillResult | None = None - working_messages: list[dict] = messages - working_system: Any | None = system + working_messages: _AnthropicMessages = messages + working_system: _AnthropicSystem = system else: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) working_messages = history_result.messages if history_result is not None else messages @@ -136,7 +144,7 @@ async def _prepare_context_managed_request( # to non-Anthropic backends that would reject them. if polyfill_will_run and history_result is None: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) return history_result @@ -144,7 +152,7 @@ async def _prepare_context_managed_request( def _polyfill_will_run( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when ``compact_20260112`` will run via the polyfill dispatcher. @@ -171,7 +179,7 @@ def _polyfill_will_run( def _spec_has_non_compact_edits( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when the spec includes edits other than ``compact_20260112``. @@ -209,9 +217,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N def _normalize_spec_edits( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run. Delegates spec-shape normalization to the dispatcher's ``_normalize_spec`` @@ -236,11 +244,11 @@ def _normalize_spec_edits( async def _run_polyfill_if_enabled( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, + messages: _AnthropicMessages, + tools: list[dict[str, object]] | None, + system: _AnthropicSystem, + context_management_spec: _ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, llm_router: "Router | None", user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -306,7 +314,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _route_openai_thinking_to_responses_api_if_needed( completion_kwargs: dict[str, Any], *, - thinking: dict[str, Any] | None, + thinking: Mapping[str, object] | None, ) -> None: """ When users call `litellm.anthropic.messages.*` with a non-Anthropic model and @@ -407,12 +415,12 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _prepare_completion_kwargs( *, max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, - system: str | list[dict[str, Any]] | None = None, + system: _AnthropicSystem = None, temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, @@ -420,7 +428,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, - extra_kwargs: dict[str, Any] | None = None, + extra_kwargs: Mapping[str, object] | None = None, ) -> tuple[dict[str, Any], dict[str, str]]: """Prepare kwargs for litellm.completion/acompletion. @@ -433,7 +441,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) - request_data: Final = { + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -528,7 +536,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, @@ -537,7 +545,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, @@ -551,10 +559,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: requested_router if requested_router is not None else _proxy_router_fallback() ) - proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( - proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None - ) + proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) polyfill_result: Final = await _prepare_context_managed_request( model=model, @@ -618,7 +623,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, @@ -627,7 +632,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, @@ -688,10 +693,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if context_management is None and not _messages_have_compaction_block(messages): polyfill_result: PolyfillResult | None = None else: - proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( - proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None - ) + proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) polyfill_result = run_async_function( _prepare_context_managed_request, model=model, diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index 671e4633af4..f7b419405ac 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -1,8 +1,9 @@ from collections.abc import Coroutine, Iterable -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypedDict import httpx from openai import AsyncAzureOpenAI, AzureOpenAI +from openai.types.shared_params.metadata import Metadata from typing_extensions import overload from ...types.llms.openai import ( @@ -22,6 +23,16 @@ from ...types.llms.openai import ( from .common_utils import BaseAzureLLM +class _RunThreadStreamData(TypedDict): + thread_id: str + assistant_id: str + additional_instructions: str | None + instructions: str | None + metadata: Metadata | None + model: str | None + tools: Iterable[AssistantToolParam] | None + + class AzureAssistantsAPI(BaseAzureLLM): def __init__(self) -> None: super().__init__() @@ -212,9 +223,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj # fmt: off @@ -301,9 +312,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj async def async_get_messages( @@ -443,7 +454,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = await openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) # fmt: off @@ -539,7 +550,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = azure_openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) async def async_get_thread( self, @@ -566,7 +577,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # fmt: off @@ -642,7 +653,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # def delete_thread(self): # pass @@ -730,7 +741,8 @@ class AzureAssistantsAPI(BaseAzureLLM): event_handler: AssistantEventHandler | None, litellm_params: dict | None = None, ) -> AssistantStreamManager[AssistantEventHandler]: - 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, @@ -740,8 +752,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) # fmt: off diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 76618e0f742..e285feb77ee 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -109,7 +109,7 @@ if MCP_AVAILABLE: ############ MCP Server REST API Routes ################# async def _safe_fire_mcp_tool_call_logging( logging_obj: Any | None, - result: Any, + result: "CallToolResult", start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1c6ad84ddb4..49a1f1314f0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -13,9 +13,9 @@ import time import traceback import types import uuid -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from fastapi import FastAPI, HTTPException @@ -145,7 +145,7 @@ try: ) # Robust auth lookup keyed by session_object. - _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() + _session_obj_auth_storage: "weakref.WeakKeyDictionary[object, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() except ImportError as e: verbose_logger.debug("MCP module not found: %s", e) MCP_AVAILABLE = False @@ -493,14 +493,14 @@ if MCP_AVAILABLE: def _gateway_create_initialization_options( self, notification_options: NotificationOptions | None = None, - experimental_capabilities: dict[str, dict[str, Any]] | None = None, + experimental_capabilities: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: opts: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) - updates: Final[dict[str, Any]] = {} + updates: Final[dict[str, str]] = {} merged: Final = _mcp_gateway_initialize_instructions.get() if merged is not None: updates["instructions"] = merged @@ -549,6 +549,17 @@ if MCP_AVAILABLE: _stateful_session_locks: Final[dict[str, asyncio.Lock]] = {} _stateful_session_active_request_counts: Final[dict[str, int]] = {} + class _TerminableTransport(Protocol): + async def terminate(self) -> None: ... + + class _TransportRegistry(Protocol): + def __contains__(self, session_id: object, /) -> bool: ... + + def pop(self, session_id: str, default: None, /) -> "_TerminableTransport | None": ... + + def _stateful_server_instances() -> _TransportRegistry: + return getattr(session_manager_stateful, "_server_instances", {}) + def _remove_stateful_session_tracking(session_id: str) -> None: _stateful_session_auth_contexts.pop(session_id, None) _stateful_session_auth_context_last_seen.pop(session_id, None) @@ -578,8 +589,8 @@ if MCP_AVAILABLE: ) -> None: """Terminate expired stateful sessions and drop their auth contexts.""" now = time.monotonic() if now is None else now - server_instances: Final = getattr(session_manager_stateful, "_server_instances", {}) - expired_session_ids: Final = [] + server_instances: Final = _stateful_server_instances() + expired_session_ids: Final[list[str]] = [] for session_id, last_seen in _stateful_session_auth_context_last_seen.items(): if _stateful_session_active_request_counts.get(session_id, 0) > 0: continue @@ -619,7 +630,7 @@ if MCP_AVAILABLE: session may proceed, or ``False`` when the caller is already at the cap with every session in flight (the new ``initialize`` should be rejected). """ - server_instances: Final = getattr(session_manager_stateful, "_server_instances", {}) + server_instances: Final = _stateful_server_instances() def _owned_live_session_ids() -> list[str]: return [ @@ -778,7 +789,7 @@ if MCP_AVAILABLE: get_virtual_tool_definitions, ) - return [Tool(**d) for d in get_virtual_tool_definitions()] + return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") @@ -847,7 +858,7 @@ if MCP_AVAILABLE: async def _build_virtual_call_logging_obj( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth, ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual @@ -885,7 +896,7 @@ if MCP_AVAILABLE: async def _dispatch_virtual_mcp_tool( name: str, - arguments: dict[str, Any] | None, + arguments: dict[str, object] | None, user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, mcp_servers: list[str] | None = None, @@ -957,7 +968,7 @@ if MCP_AVAILABLE: ) @server.call_tool() - async def mcp_server_tool_call(name: str, arguments: dict[str, Any] | None) -> CallToolResult: + async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult: """ Call a specific tool with the provided arguments Args: @@ -1621,7 +1632,7 @@ if MCP_AVAILABLE: async def _get_user_oauth_extra_headers_from_db( server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, - prefetched_creds: dict[str, dict[str, Any]] | None = None, + prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None, ) -> dict[str, str] | None: """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. @@ -1646,7 +1657,7 @@ if MCP_AVAILABLE: Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. """ - user_id: Final = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None if not user_id: return {} try: @@ -1871,7 +1882,7 @@ if MCP_AVAILABLE: list_tools_start_time: Final = datetime.now() litellm_logging_obj: LiteLLMLoggingObj | None = None - list_tools_request_data: dict[str, Any] = {} + list_tools_request_data: dict[str, object] = {} if log_list_tools_to_spendlogs: # This is intentionally minimal: only async_success_handler / post_call_failure_hook @@ -1879,7 +1890,7 @@ if MCP_AVAILABLE: list_tools_call_id: Final = str(uuid.uuid4()) # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers) - spend_logs_metadata: Final[dict[str, Any]] = { + spend_logs_metadata: Final[dict[str, object]] = { "mcp_operation": "list_tools", } if isinstance(list_tools_log_source, str): @@ -2615,7 +2626,7 @@ if MCP_AVAILABLE: async def execute_mcp_tool( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], allowed_mcp_servers: list[MCPServer], start_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -2882,7 +2893,7 @@ if MCP_AVAILABLE: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) _request_resolved_auth_headers.reset(_resolved_token) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) # Try managed MCP server tool (the name is bare; the prefix boundary was # already resolved above against this server's registered prefixes) @@ -2956,7 +2967,7 @@ if MCP_AVAILABLE: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args local_content = await _handle_local_mcp_tool(original_tool_name, arguments) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) return await _run_post_mcp_call_guardrails( result=response, @@ -3003,7 +3014,7 @@ if MCP_AVAILABLE: async def _fire_mcp_tool_call_logging( logging_obj: LiteLLMLoggingObj, - result: Any, + result: CallToolResult, start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -3070,7 +3081,7 @@ if MCP_AVAILABLE: @client async def call_mcp_tool( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3161,7 +3172,7 @@ if MCP_AVAILABLE: async def mcp_get_prompt( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3262,7 +3273,7 @@ if MCP_AVAILABLE: def _get_standard_logging_mcp_tool_call( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], server_name: str | None, session_id: str | None = None, ) -> StandardLoggingMCPToolCall: @@ -3291,13 +3302,13 @@ if MCP_AVAILABLE: async def _handle_managed_mcp_tool( server_name: str, name: str, - arguments: dict[str, Any], + arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, - litellm_logging_obj: Any | None = None, + litellm_logging_obj: LiteLLMLoggingObj | None = None, host_progress_callback: Callable | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" @@ -3320,7 +3331,7 @@ if MCP_AVAILABLE: return call_tool_result async def _handle_local_mcp_tool( - name: str, arguments: dict[str, Any] + name: str, arguments: dict[str, object] ) -> list[TextContent | ImageContent | EmbeddedResource]: """ Handle tool execution for local registry tools @@ -3426,7 +3437,8 @@ if MCP_AVAILABLE: Extract mcp-session-id from ASGI scope headers. Returns None if not present. """ - for header_name, header_value in scope.get("headers", []): + scope_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) + for header_name, header_value in scope_headers: name = header_name if isinstance(header_name, bytes) else header_name.encode() if name.lower() == b"mcp-session-id": return header_value.decode() if isinstance(header_value, bytes) else str(header_value) @@ -3528,7 +3540,7 @@ if MCP_AVAILABLE: if message.get("type") != "http.request": break - body = message.get("body", b"") or b"" + body: bytes = message.get("body", b"") or b"" if body: # Only retain up to the remaining peek budget for sniffing. # The full ``message`` is already in memory (delivered by @@ -3571,9 +3583,9 @@ if MCP_AVAILABLE: Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header: Final = b"mcp-session-id" - _headers: Final = scope.get("headers", []) + _headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) - def _normalize_header_name(header_name: Any) -> bytes | None: + def _normalize_header_name(header_name: object) -> bytes | None: if isinstance(header_name, bytes): return header_name.lower() if isinstance(header_name, str): @@ -3902,7 +3914,8 @@ if MCP_AVAILABLE: def _get_authorization_header_from_scope(scope: Scope) -> str | None: """First ``Authorization`` header value in the ASGI scope, or None.""" - for key, value in scope.get("headers", []): + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + for key, value in scope_headers: if key.lower() == b"authorization": return value.decode("latin-1") return None @@ -3921,7 +3934,8 @@ if MCP_AVAILABLE: ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream would leak the proxy key to a third-party MCP server. """ - has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", [])) + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope_headers) if not has_litellm_key_header: return None return _get_authorization_header_from_scope(scope) @@ -4115,7 +4129,7 @@ if MCP_AVAILABLE: async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, @@ -4135,7 +4149,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar (set by # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py). @@ -4436,7 +4451,7 @@ if MCP_AVAILABLE: async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through SSE.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, @@ -4456,7 +4471,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar so the # downstream probe list matches the fully-authorized server set @@ -4680,7 +4696,8 @@ if MCP_AVAILABLE: ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": - for key, value in message.get("headers", []): + response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", []) + for key, value in response_headers: header_name = key if isinstance(key, bytes) else str(key).encode() if header_name.lower() == b"mcp-session-id": session_id = value.decode() if isinstance(value, bytes) else str(value) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 761d8aabc8a..b68d4a68b79 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -6,10 +6,10 @@ import concurrent.futures import inspect import json import os -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone from types import UnionType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, Union, cast, get_args, get_origin +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request @@ -54,8 +54,8 @@ from litellm.types.guardrails import ( if TYPE_CHECKING: from types import CodeType - from prisma.actions import LiteLLM_GuardrailsTableActions from prisma.models import LiteLLM_GuardrailsTable + from pydantic.fields import FieldInfo from litellm.proxy.utils import PrismaClient @@ -65,24 +65,44 @@ router: Final = APIRouter() GUARDRAIL_REGISTRY: Final = GuardrailRegistry() -def _guardrails_table(prisma_client: "PrismaClient") -> "LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]": - table: Final[LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]] = GuardrailsRepository(prisma_client).table +class _GuardrailsTableActions(Protocol): + async def create(self, data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": ... + + async def delete(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ... + + async def find_unique(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ... + + async def find_many( + self, where: Mapping[str, object], order: Mapping[str, str] + ) -> "Sequence[LiteLLM_GuardrailsTable]": ... + + async def update( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "LiteLLM_GuardrailsTable | None": ... + + +def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]: + return mapping + + +def _guardrails_table(prisma_client: "PrismaClient") -> _GuardrailsTableActions: + table: Final[_GuardrailsTableActions] = GuardrailsRepository(prisma_client).table return table async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": - row: Final[LiteLLM_GuardrailsTable] = await GuardrailsRepository(prisma_client).table.create(data=data) + row: Final = await _guardrails_table(prisma_client).create(data=data) return row async def _delete_guardrail_row(prisma_client: "PrismaClient", where: Mapping[str, object]) -> None: - await GuardrailsRepository(prisma_client).table.delete(where=where) + await _guardrails_table(prisma_client).delete(where=where) async def _find_team_guardrail_rows( prisma_client: "PrismaClient", where: Mapping[str, object] ) -> "Sequence[LiteLLM_GuardrailsTable]": - rows: Final[Sequence[LiteLLM_GuardrailsTable]] = await GuardrailsRepository(prisma_client).table.find_many( + rows: Final = await _guardrails_table(prisma_client).find_many( where=where, order={"created_at": "desc"}, ) @@ -499,10 +519,12 @@ async def update_guardrail( if existing_guardrail is None: raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found") - result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db( - guardrail_id=guardrail_id, - guardrail=request.guardrail, - prisma_client=prisma_client, + result: Final = _as_str_object_mapping( + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=request.guardrail, + prisma_client=prisma_client, + ) ) guardrail_name: Final = result.get("guardrail_name", "Unknown") @@ -613,7 +635,7 @@ class RegisterGuardrailRequest(BaseModel): """Request body for POST /guardrails/register. Follows Generic Guardrail API config.""" guardrail_name: str - litellm_params: dict[str, Any] # guardrail, mode, api_base required; api_key, headers, etc. optional + litellm_params: dict[str, object] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: dict[str, object] | None = None team_id: str | None = None @@ -1172,12 +1194,14 @@ async def patch_guardrail( ) # Update litellm_params if default_on is provided or pii_entities_config is provided - litellm_params = LitellmParams(**dict(existing_guardrail.get("litellm_params", {}))) + existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {}))) + litellm_params = LitellmParams(**existing_litellm_params) if request.litellm_params is not None: requested_litellm_params: Final = request.litellm_params.model_dump(exclude_unset=True) litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True) litellm_params_dict.update(requested_litellm_params) - litellm_params = LitellmParams(**litellm_params_dict) + merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict) + litellm_params = LitellmParams(**merged_litellm_params) # Update guardrail_info if provided guardrail_info: Final = ( @@ -1193,10 +1217,12 @@ async def patch_guardrail( litellm_params=litellm_params, guardrail_info=guardrail_info, ) - result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db( - guardrail_id=guardrail_id, - guardrail=guardrail, - prisma_client=prisma_client, + result: Final = _as_str_object_mapping( + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=guardrail, + prisma_client=prisma_client, + ) ) guardrail_name = result.get("guardrail_name", "Unknown") @@ -1552,31 +1578,46 @@ async def validate_blocked_words_file(request: dict[str, str]): return {"valid": False, "error": f"Validation error: {e}"} -def _get_field_type_from_annotation(field_annotation: Any) -> str: +def _dunder_origin(annotation: object) -> object: + origin: Final[object] = getattr(annotation, "__origin__", None) + return origin + + +def _dunder_name(annotation: object) -> object: + name: Final[object] = getattr(annotation, "__name__", None) + return name + + +def _dunder_args(annotation: object) -> tuple[object, ...]: + args: Final[tuple[object, ...]] = getattr(annotation, "__args__", ()) + return args + + +def _get_field_type_from_annotation(field_annotation: object) -> str: """ Convert a Python type annotation to a UI-friendly type string """ # Handle Union types (like Optional[T]) if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType: # For Optional[T], get the non-None type - args: Final = get_args(field_annotation) + args: Final[tuple[object, ...]] = get_args(field_annotation) non_none_args: Final = [arg for arg in args if arg is not type(None)] if non_none_args: field_annotation = non_none_args[0] # Handle List types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is list: + if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is list: return "array" # Handle Dict types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is dict: + if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is dict: return "dict" # Handle Literal types if hasattr(field_annotation, "__origin__") and hasattr(field_annotation, "__args__"): # Check for Literal types (Python 3.8+) - origin: Final = field_annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": + origin: Final = _dunder_origin(field_annotation) + if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal": return "select" # For dropdown/select inputs # Handle basic types @@ -1595,66 +1636,66 @@ def _get_field_type_from_annotation(field_annotation: Any) -> str: return "string" -def _extract_literal_values(annotation: Any) -> list[str]: +def _extract_literal_values(annotation: object) -> Sequence[object]: """ Extract literal values from a Literal type annotation """ if hasattr(annotation, "__origin__") and hasattr(annotation, "__args__"): - origin: Final = annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": - return list(annotation.__args__) + origin: Final = _dunder_origin(annotation) + if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal": + return list(_dunder_args(annotation)) return [] -def _get_dict_key_options(field_annotation: Any) -> list[str] | None: +def _get_dict_key_options(field_annotation: object) -> Sequence[object] | None: """ Extract key options from Dict[Literal[...], T] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict + and _dunder_origin(field_annotation) is dict and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 2: key_type: Final = args[0] return _extract_literal_values(key_type) return None -def _get_dict_value_type(field_annotation: Any) -> str: +def _get_dict_value_type(field_annotation: object) -> str: """ Get the value type from Dict[K, V] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict + and _dunder_origin(field_annotation) is dict and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 2: value_type: Final = args[1] return _get_field_type_from_annotation(value_type) return "string" -def _get_list_element_options(field_annotation: Any) -> list[str] | None: +def _get_list_element_options(field_annotation: object) -> Sequence[object] | None: """ Extract element options from List[Literal[...]] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is list + and _dunder_origin(field_annotation) is list and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 1: element_type: Final = args[0] return _extract_literal_values(element_type) return None -def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool: +def _should_skip_optional_params(field_name: str, field_annotation: object) -> bool: """Check if optional_params field should be skipped (not meaningfully overridden).""" if field_name != "optional_params": return False @@ -1664,12 +1705,12 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Check if the annotation is still a generic TypeVar (not specialized) if isinstance(field_annotation, TypeVar) or ( - hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar + hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is TypeVar ): return True # Also skip if it's a generic type that wasn't specialized - if hasattr(field_annotation, "__name__") and field_annotation.__name__ in ( + if hasattr(field_annotation, "__name__") and _dunder_name(field_annotation) in ( "T", "TypeVar", ): @@ -1677,18 +1718,18 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Handle Optional[T] where T is still a TypeVar if hasattr(field_annotation, "__args__"): - non_none_args: Final = [arg for arg in field_annotation.__args__ if arg is not type(None)] + non_none_args: Final = [arg for arg in _dunder_args(field_annotation) if arg is not type(None)] if non_none_args and isinstance(non_none_args[0], TypeVar): return True return False -def _unwrap_optional_type(field_annotation: Any) -> Any: +def _unwrap_optional_type(field_annotation: object) -> object: """Unwrap Optional types to get the actual type.""" if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType: # For Optional[BaseModel], get the non-None type - args: Final = get_args(field_annotation) + args: Final[tuple[object, ...]] = get_args(field_annotation) non_none_args: Final = [arg for arg in args if arg is not type(None)] if non_none_args: return non_none_args[0] @@ -1696,20 +1737,20 @@ def _unwrap_optional_type(field_annotation: Any) -> Any: def _build_field_dict( - field: Any, - field_annotation: Any, + field: "FieldInfo", + field_annotation: object, description: str, required: bool, -) -> dict[str, Any]: +) -> dict[str, object]: """Build field dictionary for non-nested fields.""" # Determine the field type from annotation field_type = _get_field_type_from_annotation(field_annotation) # Check for custom UI type override - field_json_schema_extra: Final = getattr(field, "json_schema_extra", {}) + field_json_schema_extra: Final[Mapping[str, object]] = getattr(field, "json_schema_extra", {}) if field_json_schema_extra and "ui_type" in field_json_schema_extra: ui_type: Final = field_json_schema_extra["ui_type"] - field_type = ui_type.value if hasattr(ui_type, "value") else ui_type + field_type = getattr(ui_type, "value", ui_type) elif field_json_schema_extra and "type" in field_json_schema_extra: field_type = field_json_schema_extra["type"] @@ -1748,8 +1789,9 @@ def _build_field_dict( field_dict["options"] = literal_options # Add default value if it exists - if field.default is not None and field.default is not ...: - field_dict["default_value"] = field.default + field_default: Final[object] = getattr(field, "default", None) + if field_default is not None and field_default is not ...: + field_dict["default_value"] = field_default # Copy min, max, step from json_schema_extra for number/percentage inputs if field_json_schema_extra: @@ -1763,7 +1805,7 @@ def _build_field_dict( def _extract_fields_recursive( model: type[BaseModel], depth: int = 0, -) -> dict[str, Any]: +) -> dict[str, object]: # Check if we've exceeded the maximum recursion depth if depth > DEFAULT_MAX_RECURSE_DEPTH: raise HTTPException( @@ -1817,7 +1859,7 @@ def _extract_fields_recursive( return fields -def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, Any]: +def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, object]: """ Get the fields from a Pydantic model as a nested dictionary structure """ @@ -2141,7 +2183,26 @@ def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: return "response" if input_type == "response" else "request" -def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGuardrailRequest) -> None: +class _GuardrailLoggingObj(Protocol): + call_type: str + model_call_details: dict[str, object] + + @property + def update_messages(self) -> "Callable[..., object]": ... + + @property + def async_success_handler(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def success_handler(self) -> "Callable[..., object]": ... + + +class _GuardrailProxyLogging(Protocol): + @property + def post_call_success_hook(self) -> "Callable[..., Awaitable[object]]": ... + + +def _patch_logging_obj_for_guardrail(litellm_logging_obj: _GuardrailLoggingObj, request: ApplyGuardrailRequest) -> None: """Configure the logging object so Langfuse/OTEL extract input and output correctly.""" litellm_logging_obj.call_type = "pass_through_endpoint" litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint" @@ -2151,8 +2212,8 @@ def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGua async def _emit_guardrail_success_logs( - proxy_logging_obj: Any, - litellm_logging_obj: Any, + proxy_logging_obj: _GuardrailProxyLogging, + litellm_logging_obj: _GuardrailLoggingObj | None, data: dict, user_api_key_dict: UserAPIKeyAuth, response: ApplyGuardrailResponse, 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 068a3ecf31b..facb822d00d 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 @@ -19,7 +19,7 @@ request is sent with the ``X-Cisco-AI-Defense-API-Key`` header. import json import os -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -94,13 +94,13 @@ class _CiscoVerdict: is_safe: bool | None classifications: list[str] severity: str | None - rules: list[dict[str, Any]] + rules: list[dict[str, object]] explanation: str | None event_id: str | None action: str | None = None sanitized_text: str | None = None - sanitized_messages: list[dict[str, Any]] | None = None - sanitized_mcp_arguments: dict[str, Any] | None = None + sanitized_messages: list[dict[str, object]] | None = None + sanitized_mcp_arguments: dict[str, object] | None = None class CiscoAIDefenseGuardrailMissingSecrets(Exception): @@ -136,7 +136,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): api_base: str | None = None, inspection_type: str | None = None, inspect_path: str | None = None, - enabled_rules: list[dict[str, Any]] | None = None, + enabled_rules: Sequence[object] | None = None, integration_profile_id: str | None = None, integration_profile_version: str | None = None, integration_tenant_id: str | None = None, @@ -415,7 +415,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: AsyncIterator[Any], + response: AsyncIterator[object], request_data: dict, ): """Buffer and inspect streaming chat output before delivery.""" @@ -437,7 +437,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self.guardrail_name, ) - all_chunks: Final[list[Any]] = [] + all_chunks: Final[list[object]] = [] try: async for chunk in response: all_chunks.append(chunk) @@ -497,7 +497,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): response_obj=assembled, ) except HTTPException as exc: - error_obj: dict[str, Any] = self._http_exception_to_error_obj(exc) + error_obj: dict[str, object] = self._http_exception_to_error_obj(exc) verbose_proxy_logger.warning( "Cisco AI Defense guardrail (%s): streaming response " "blocked — emitting SSE error event instead of " @@ -531,7 +531,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): for chunk in all_chunks: yield chunk - def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, Any]: + def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, object]: """Canonical block payload used across all four block paths. Same dict is the ``HTTPException.detail`` for chat / MCP request @@ -555,34 +555,34 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): "event_id": verdict.event_id, } - def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, Any]: + def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, object]: """Wrap an ``HTTPException`` detail into the SSE ``error`` payload. For Cisco's own blocks the detail is already the canonical block payload, so this is a near-passthrough that just adds ``code`` / ``guardrail`` defaults for non-Cisco / unstructured details. """ - error_obj: dict[str, Any] = dict(exc.detail) if isinstance(exc.detail, dict) else {"message": str(exc.detail)} + error_obj: dict[str, object] = {**exc.detail} if isinstance(exc.detail, dict) else {"message": str(exc.detail)} error_obj.setdefault("message", error_obj.get("error", "Guardrail block")) error_obj.setdefault("code", exc.status_code) error_obj.setdefault("guardrail", self.guardrail_name) return error_obj @classmethod - def _streaming_content_was_modified(cls, original_chunks: list[Any], assembled: ModelResponse) -> bool: + def _streaming_content_was_modified(cls, original_chunks: Sequence[object], assembled: ModelResponse) -> bool: """Decide whether redact changed content or tool/function arguments.""" original_text: Final = cls._extract_streaming_chunk_scan_text(original_chunks) assembled_text: Final = " ".join(m.get("content", "") for m in cls._extract_response_messages(assembled)) return original_text != assembled_text @classmethod - def _extract_streaming_chunk_scan_text(cls, chunks: list[Any]) -> str: + def _extract_streaming_chunk_scan_text(cls, chunks: Sequence[object]) -> str: original_text = "" argument_text = "" for chunk in chunks: choices = getattr(chunk, "choices", None) or [] for c in choices: - delta = getattr(c, "delta", None) + delta: object | None = getattr(c, "delta", None) if delta is None: continue text = getattr(delta, "content", None) @@ -595,7 +595,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): args = cls._extract_tool_call_arguments(tc) if args: argument_text += args - fc = getattr(delta, "function_call", None) + fc: object | None = getattr(delta, "function_call", None) if fc is not None: args = cls._extract_function_call_arguments(fc) if args: @@ -673,7 +673,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): allow, WARNING for intervened/redacted, ERROR is left for upstream API failures. """ - fields: Final[dict[str, Any]] = { + fields: Final[dict[str, object]] = { "guardrail": self.guardrail_name, "surface": context.surface, "direction": context.direction, @@ -752,7 +752,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, direction: str = "input", response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: url: Final = f"{self.api_base}{self.inspect_path}" payload: Final = self._build_chat_payload(messages, request_data, user_api_key_dict) start_time: Final = datetime.now() @@ -784,7 +784,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): messages: list[dict[str, str]], request_data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: + ) -> dict[str, object]: return { "messages": messages, "metadata": self._build_metadata(request_data, user_api_key_dict), @@ -798,9 +798,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): async def _post_inspection( self, url: str, - payload: dict[str, Any], + payload: dict[str, object], surface: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: headers: Final = self._build_headers() verbose_proxy_logger.debug( "Cisco AI Defense guardrail: posting %s inspection to %s", @@ -856,8 +856,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: - metadata: Final[dict[str, Any]] = {} + ) -> dict[str, object]: + metadata: Final[dict[str, object]] = {} user: Final = request_data.get("user") or getattr(user_api_key_dict, "user_id", None) if user: @@ -884,8 +884,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return metadata - def _build_config(self) -> dict[str, Any]: - config: Final[dict[str, Any]] = {} + def _build_config(self) -> dict[str, object]: + config: Final[dict[str, object]] = {} if self.enabled_rules: config["enabled_rules"] = self.enabled_rules if self.integration_profile_id: @@ -899,7 +899,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return config @staticmethod - def _normalize_rule(rule: object) -> dict[str, Any]: + def _normalize_rule(rule: object) -> dict[str, object]: """Coerce a user-supplied rule into the wire-shape dict Cisco expects. Accepts ``str``, ``dict``, and Pydantic model inputs. @@ -922,7 +922,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): rule = dumped if isinstance(rule, dict): - normalized: Final[dict[str, Any]] = {} + normalized: Final[dict[str, object]] = {} rule_name: Final = rule.get("rule_name") if rule_name: normalized["rule_name"] = rule_name @@ -950,7 +950,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): context: _ScanContext, start_time: datetime, response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Parse, log, and (optionally) raise/redact on the Cisco verdict. ``context.direction`` is ``"input"`` for request scans and ``"output"`` @@ -1119,10 +1119,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @classmethod def _sanitize_response_for_logging( cls, - inspect_response: dict[str, Any], + inspect_response: Mapping[str, object], surface: str, action: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Drop bulky / privacy-sensitive fields, recursing into nested dicts. MCP verdicts are commonly nested under ``result``, so a @@ -1138,9 +1138,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return sanitized @classmethod - def _strip_sensitive_keys(cls, d: dict[str, Any]) -> dict[str, Any]: + def _strip_sensitive_keys(cls, d: Mapping[str, object]) -> dict[str, object]: """Recursively strip privacy-sensitive keys from a verdict dict.""" - out: Final[dict[str, Any]] = {} + out: Final[dict[str, object]] = {} for key, value in d.items(): if key.startswith("_") or key in cls._REDACTED_LOG_KEYS: continue @@ -1222,8 +1222,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_jsonrpc_error( - inspect_response: dict[str, Any], - ) -> dict[str, Any] | None: + inspect_response: Mapping[str, object], + ) -> dict[str, object] | None: """Detect a JSON-RPC error envelope inside an HTTP 200 response. The Cisco Inspect API can return ``{"error": {...}}`` (or nest one @@ -1270,7 +1270,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_sanitized_text( - inspect_response: dict[str, Any], + inspect_response: Mapping[str, object], ) -> str | None: """Pull ``sanitized_text`` (or camelCase variant) off the verdict.""" for key in ("sanitized_text", "sanitizedText"): @@ -1287,8 +1287,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_sanitized_messages( - inspect_response: dict[str, Any], - ) -> list[dict[str, Any]] | None: + inspect_response: Mapping[str, object], + ) -> list[dict[str, object]] | None: """Pull a sanitized OpenAI-format messages array off the verdict. Cisco can return the rewrite under several keys; we accept any of @@ -1354,7 +1354,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_mcp_input( request_data: dict, sanitized_text: str | None, - sanitized_mcp_arguments: dict[str, Any] | None, + sanitized_mcp_arguments: dict[str, object] | None, ) -> bool: """Rewrite MCP request arguments in all locations the proxy reads.""" if sanitized_mcp_arguments is not None: @@ -1388,7 +1388,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, request_data: dict, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite chat request input (``messages`` or ``input``).""" if sanitized_messages and self._extract_tool_definition_text(request_data): @@ -1444,7 +1444,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): cls, request_data: dict, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: if sanitized_messages: instruction_text: Final = cls._instruction_text_from_messages(sanitized_messages) @@ -1457,7 +1457,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return False @classmethod - def _instruction_text_from_messages(cls, messages: list[dict[str, Any]]) -> str | None: + def _instruction_text_from_messages(cls, messages: list[dict[str, object]]) -> str | None: for message in messages: if not isinstance(message, dict): continue @@ -1468,7 +1468,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return None @classmethod - def _non_instruction_messages(cls, messages: list[dict[str, Any]] | None) -> list[dict[str, Any]] | None: + def _non_instruction_messages(cls, messages: list[dict[str, object]] | None) -> list[dict[str, object]] | None: if messages is None: return None return [ @@ -1499,7 +1499,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, response_obj: object, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite chat response (``ModelResponse`` or ``ResponsesAPIResponse``).""" if response_obj is None: @@ -1526,7 +1526,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_model_response_choices( choices: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Redact every returned choice, including tool-call/reasoning fields.""" if sanitized_messages: @@ -1570,7 +1570,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_text_completion_choices( choices: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite ``/v1/completions`` text choices after Cisco redaction.""" replacement = sanitized_text @@ -1638,7 +1638,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, output_items: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: replacement_text: str | None = sanitized_text if not replacement_text and sanitized_messages: @@ -1672,14 +1672,14 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _sanitized_messages_to_responses_input( - sanitized_messages: list[dict[str, Any]], - ) -> list[dict[str, Any]] | None: + sanitized_messages: list[dict[str, object]], + ) -> list[dict[str, object]] | None: """Convert chat-shape sanitized_messages to Responses API ``input``. Returns ``None`` if nothing usable could be converted, so the caller falls back to ``on_flagged_action``. """ - out: Final[list[dict[str, Any]]] = [] + out: Final[list[dict[str, object]]] = [] for m in sanitized_messages: if not isinstance(m, dict): continue @@ -1764,7 +1764,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): start_time: datetime | None = None, surface: str = "chat", direction: str = "input", - ) -> dict[str, Any]: + ) -> dict[str, object]: verbose_proxy_logger.error( "Cisco AI Defense guardrail (%s): API communication failed: %s", surface, @@ -2060,7 +2060,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return getattr(obj, key, None) @classmethod - def _field_list(cls, obj: object, key: str) -> list[Any]: + def _field_list(cls, obj: object, key: str) -> list[object]: value: Final = cls._field(obj, key) return value if isinstance(value, list) else [] diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index c29da89b15f..5bbb01c6c8e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -8,8 +8,8 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint import copy import json -from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, Final +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import HTTPException @@ -34,6 +34,9 @@ if TYPE_CHECKING: # Imported lazily at runtime (inside the streaming hook) to avoid a # module-level cyclic import with litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) @@ -41,12 +44,35 @@ A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME: Final = "unified_llm_guardrails" +class _EndpointTranslation(Protocol): + @property + def process_input_messages(self) -> "Callable[..., Awaitable[dict[str, object]]]": ... + + @property + def process_output_response(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... + + +def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation: + return translation + + +def _chunk_choices(item: object) -> Sequence[object]: + choices: Final[Sequence[object]] = getattr(item, "choices", None) or [] + return choices + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" -def _get_a2a_request_id(responses_so_far: list[Any], request_data: dict) -> str | None: +def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) -> str | None: """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting.""" for item in responses_so_far: if isinstance(item, dict) and "id" in item: @@ -138,7 +164,9 @@ class UnifiedLLMGuardrails(CustomLogger): except ValueError: return data # handle unmapped call types - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) _ensure_litellm_metadata(data, user_api_key_dict) @@ -156,7 +184,7 @@ class UnifiedLLMGuardrails(CustomLogger): async def async_moderation_hook( self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral - ) -> Any: + ) -> object: """ Runs in parallel to LLM API call Runs on only Input @@ -187,7 +215,9 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings: return data - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) _ensure_litellm_metadata(data, user_api_key_dict) @@ -202,7 +232,7 @@ class UnifiedLLMGuardrails(CustomLogger): data: dict, user_api_key_dict: UserAPIKeyAuth, response, - ) -> Any: + ) -> object: """ Runs on response from LLM API call @@ -271,7 +301,9 @@ class UnifiedLLMGuardrails(CustomLogger): ) return response - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) try: response = await endpoint_translation.process_output_response( @@ -299,10 +331,10 @@ class UnifiedLLMGuardrails(CustomLogger): async def _handle_streaming_block( self, exc: "ModifyResponseException", - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, stream_started: bool, - responses_so_far: list[Any], - ) -> AsyncGenerator[Any, None]: + responses_so_far: Sequence[object], + ) -> AsyncGenerator[object, None]: """ Terminate a streamed response cleanly when a guardrail blocks it. @@ -323,7 +355,7 @@ class UnifiedLLMGuardrails(CustomLogger): @staticmethod def _resolve_transform_call_type( user_api_key_dict: UserAPIKeyAuth, - mappings: dict, + mappings: Mapping[CallTypes, type["BaseTranslation"]], ) -> str | None: """Resolve the call type for the incremental_diff path, or None if the route is unresolvable / unsupported. @@ -356,9 +388,9 @@ class UnifiedLLMGuardrails(CustomLogger): self, exc: HTTPException, call_type: str | None, - responses_so_far: list[Any], + responses_so_far: Sequence[object], request_data: dict, - ) -> AsyncGenerator[Any, None]: + ) -> AsyncGenerator[object, None]: """Surface a mid-stream HTTPException. For A2A (NDJSON) call types the response has already started, so emit an in-stream JSON-RPC error chunk; otherwise re-raise so the proxy can report it. @@ -387,7 +419,7 @@ class UnifiedLLMGuardrails(CustomLogger): def _build_transform_chunk( self, *, - reference_chunk: Any, + reference_chunk: object, mutated_text_per_choice: dict[int, str], emitted_text_per_choice: dict[int, str], holdback_per_choice: dict[int, int], @@ -500,18 +532,18 @@ class UnifiedLLMGuardrails(CustomLogger): async def _emit_transform_round( self, *, - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, request_data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str, - reference_chunk: Any, - responses_so_far: list[Any], - responses_yielded: list[Any], + reference_chunk: object, + responses_so_far: Sequence[object], + responses_yielded: list[object], emitted_text_per_choice: dict[int, str], finish_reason_per_choice: dict[int, str | None], is_final: bool, - ) -> AsyncGenerator[Any, None]: + ) -> AsyncGenerator[object, None]: """Run one guardrail processing round and emit the resulting diff chunk. Raises ``_StreamTerminated`` (after emitting the terminal block message or @@ -564,14 +596,14 @@ class UnifiedLLMGuardrails(CustomLogger): self, *, guardrail_to_apply: CustomGuardrail, - response: Any, + response: AsyncIterable[object], request_data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str, sampling_rate: int, end_of_stream_only: bool, - mappings: dict, - ) -> AsyncGenerator[Any, None]: + mappings: Mapping[CallTypes, type["BaseTranslation"]], + ) -> AsyncGenerator[object, None]: """Emit guardrail text transformations as new deltas on the stream. Raw chunks are withheld and accumulated; on each sampled processing round @@ -580,15 +612,15 @@ class UnifiedLLMGuardrails(CustomLogger): synthetic chunk. A BLOCK terminates the stream via the shared block handler; an underflow surfaces as an HTTPException. """ - endpoint_translation: Final = mappings[CallTypes(call_type)]() - responses_so_far: Final[list[Any]] = [] - responses_yielded: Final[list[Any]] = [] + endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + responses_so_far: Final[list[object]] = [] + responses_yielded: Final[list[object]] = [] emitted_text_per_choice: Final[dict[int, str]] = {} finish_reason_per_choice: Final[dict[int, str | None]] = {} chunk_counter = 0 - last_chunk: Any | None = None + last_chunk: object | None = None - def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]: + def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]: return self._emit_transform_round( endpoint_translation=endpoint_translation, guardrail_to_apply=guardrail_to_apply, @@ -694,13 +726,13 @@ class UnifiedLLMGuardrails(CustomLogger): async def _inspect_full_response_for_block( self, *, - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - responses_so_far: list[Any], - responses_yielded: list[Any], - ) -> AsyncGenerator[Any, None]: + responses_so_far: Sequence[object], + responses_yielded: Sequence[object], + ) -> AsyncGenerator[object, None]: """Run the block-only guardrail inspection over the full assembled response (text + tool calls) so nothing bypasses the block decision. @@ -734,17 +766,17 @@ class UnifiedLLMGuardrails(CustomLogger): raise _StreamTerminated() @staticmethod - def _chunk_has_tool_calls(item: Any) -> bool: - for choice in getattr(item, "choices", None) or []: + def _chunk_has_tool_calls(item: object) -> bool: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) if getattr(delta, "tool_calls", None): return True return False @staticmethod - def _chunk_carries_text(item: Any) -> bool: + def _chunk_carries_text(item: object) -> bool: """True if any choice in this chunk has non-empty string ``delta.content``.""" - for choice in getattr(item, "choices", None) or []: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) content = getattr(delta, "content", None) if isinstance(content, str) and content != "": @@ -753,7 +785,7 @@ class UnifiedLLMGuardrails(CustomLogger): @staticmethod def _tool_call_passthrough_chunk( - item: Any, + item: object, finish_reason_per_choice: "dict[int, str | None] | None" = None, ) -> ModelResponseStream: """Copy of a chunk carrying tool calls with all text content stripped. @@ -772,7 +804,7 @@ class UnifiedLLMGuardrails(CustomLogger): redaction purpose. """ synthetic_choices: Final[list[StreamingChoices]] = [] - for choice in getattr(item, "choices", None) or []: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) idx = getattr(choice, "index", 0) or 0 original_finish = getattr(choice, "finish_reason", None) @@ -801,15 +833,15 @@ class UnifiedLLMGuardrails(CustomLogger): ) @staticmethod - def _record_finish_reasons(item: Any, finish_reason_per_choice: dict[int, str | None]) -> None: - for choice in getattr(item, "choices", None) or []: + def _record_finish_reasons(item: object, finish_reason_per_choice: dict[int, str | None]) -> None: + for choice in _chunk_choices(item): finish_reason = getattr(choice, "finish_reason", None) if finish_reason is not None: finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason @staticmethod - def _chunk_has_finish_reason(item: Any) -> bool: - choices: Final = getattr(item, "choices", None) or [] + def _chunk_has_finish_reason(item: object) -> bool: + choices: Final = _chunk_choices(item) return any(getattr(choice, "finish_reason", None) is not None for choice in choices) async def async_post_call_streaming_iterator_hook( @@ -845,22 +877,22 @@ class UnifiedLLMGuardrails(CustomLogger): # Get streaming configuration. Resolution order (later wins): default # < guardrail attribute < guardrail_config dict < this callback's # optional_params. - def _streaming_flag(name: str, default: Any) -> Any: + def _streaming_flag(name: str, default: object) -> Any: value = default if guardrail_to_apply is not None: value = getattr(guardrail_to_apply, name, value) - config: Final = getattr(guardrail_to_apply, "guardrail_config", {}) + config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {}) if isinstance(config, dict): value = config.get(name, value) return self.optional_params.get(name, value) - sampling_rate: Final = _streaming_flag("streaming_sampling_rate", 5) + sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). - end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False) + end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False) # "block_only" (default) drops guardrail text rewrites on the streaming # path; "incremental_diff" emits them as synthetic deltas (see # _run_incremental_transform_stream). - streaming_transform_mode: Final = _streaming_flag("streaming_transform_mode", "block_only") + streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only") # Withhold every chunk until end-of-stream moderation passes, then # release the original chunks (clean) or only the block message # (blocked) -- moderating the whole response *before* any content @@ -868,7 +900,9 @@ class UnifiedLLMGuardrails(CustomLogger): # release the original chunks are replayed as-is, so a # content-rewriting guardrail (e.g. PII masking) would leak # unredacted content. Guarded below via mask_response_content. - buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default) + buffer_until_moderated: bool = _streaming_flag( + "streaming_buffer_until_moderated", buffer_until_moderated_default + ) if ( buffer_until_moderated @@ -939,9 +973,9 @@ class UnifiedLLMGuardrails(CustomLogger): # Infer call type from first chunk call_type = None chunk_counter = 0 - responses_so_far: Final[list[Any]] = [] - responses_yielded: Final[list[Any]] = [] - pending_end_of_stream_items: Final[list[Any]] = [] + responses_so_far: Final[list[object]] = [] + responses_yielded: Final[list[object]] = [] + pending_end_of_stream_items: Final[list[object]] = [] # Whether any real response chunk has been forwarded to the client. # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index dd61cad15a1..a64ed764a67 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -26,7 +26,8 @@ Usage: import base64 import json -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -36,7 +37,10 @@ from litellm.llms.litellm_proxy.skills.prompt_injection import ( SkillPromptInjectionHandler, ) from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth -from litellm.types.utils import CallTypes, CallTypesLiteral +from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes + +if TYPE_CHECKING: + from litellm.llms.litellm_proxy.skills.sandbox_executor import SkillsSandboxExecutor class SkillsInjectionHook(CustomLogger): @@ -99,7 +103,7 @@ class SkillsInjectionHook(CustomLogger): verbose_proxy_logger.debug("SkillsInjectionHook: Processing %s skills", len(skills)) litellm_skills: Final[list[LiteLLM_SkillsTable]] = [] - anthropic_skills: Final[list[dict[str, Any]]] = [] + anthropic_skills: Final[list[dict[str, object]]] = [] # Separate skills by prefix for skill in skills: @@ -324,9 +328,9 @@ class SkillsInjectionHook(CustomLogger): async def async_post_call_success_deployment_hook( self, request_data: dict, - response: Any, + response: LLMResponseTypes, call_type: CallTypes | None, - ) -> Any | None: + ) -> LLMResponseTypes | None: """ Post-call hook to handle automatic code execution. @@ -372,7 +376,7 @@ class SkillsInjectionHook(CustomLogger): # Check if any tool call needs execution (litellm_code_execution or skill tool) has_executable_tool = False for tc in tool_calls: - tool_name = tc.get("name", "") + tool_name: str = tc.get("name", "") # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx) if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value or tool_name.startswith(LITELLM_SKILL_ID_PREFIX): has_executable_tool = True @@ -441,7 +445,7 @@ class SkillsInjectionHook(CustomLogger): data: dict, response: Any, skill_files: dict[str, bytes], - ) -> Any: + ) -> LLMResponseTypes | None: """ Execute the code execution loop for messages API (Anthropic format). @@ -466,7 +470,7 @@ class SkillsInjectionHook(CustomLogger): max_tokens: Final = data.get("max_tokens", 4096) executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[dict[str, object]]] = [] current_response = response for iteration in range(self.max_iterations): @@ -511,9 +515,9 @@ class SkillsInjectionHook(CustomLogger): # Process tool calls tool_results = [] for tc in tool_calls: - tool_name = tc.get("name", "") + tool_name: str = tc.get("name", "") tool_id = tc.get("id", "") - tool_input = tc.get("input", {}) + tool_input: Mapping[str, str] = tc.get("input", {}) # Execute if it's litellm_code_execution OR a skill tool if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: @@ -561,8 +565,8 @@ class SkillsInjectionHook(CustomLogger): self, code: str, skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute code in sandbox and return result string.""" try: @@ -574,7 +578,8 @@ class SkillsInjectionHook(CustomLogger): # Collect generated files if exec_result.get("files"): - for f in exec_result["files"]: + files: Final[Sequence[Mapping[str, str]]] = exec_result["files"] + for f in files: generated_files.append( { "name": f["name"], @@ -595,10 +600,10 @@ class SkillsInjectionHook(CustomLogger): async def _execute_skill_tool( self, tool_name: str, - tool_input: dict[str, Any], + tool_input: Mapping[str, str], skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute a skill tool by generating and running code based on skill content.""" # Generate code based on available skill modules @@ -670,7 +675,7 @@ print('No executable skill module found') data: dict, response: Any, skill_files: dict[str, bytes], - ) -> Any: + ) -> LLMResponseTypes: """ Execute the code execution loop until model gives final response. @@ -704,7 +709,7 @@ print('No executable skill module found') kwargs: Final = {k: v for k, v in data.items() if k not in _EXCLUDED_ACOMPLETION_KEYS} executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[dict[str, object]]] = [] current_response: Any = response for iteration in range(self.max_iterations): @@ -713,7 +718,7 @@ print('No executable skill module found') stop_reason = current_response.choices[0].finish_reason # Build assistant message for conversation history - assistant_msg_dict: dict[str, Any] = { + assistant_msg_dict: dict[str, object] = { "role": "assistant", "content": assistant_message.content, } @@ -781,13 +786,13 @@ print('No executable skill module found') self, tool_call: Any, skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute a litellm_code_execution tool call and return result string.""" try: args: Final = json.loads(tool_call.function.arguments) - code: Final = args.get("code", "") + code: Final[str] = args.get("code", "") verbose_proxy_logger.debug("SkillsInjectionHook: Executing code (%s chars)", len(code)) @@ -802,7 +807,8 @@ print('No executable skill module found') # Collect generated files if exec_result.get("files"): tool_result += "\n\nGenerated files:" - for f in exec_result["files"]: + files: Final[Sequence[Mapping[str, str]]] = exec_result["files"] + for f in files: file_content = base64.b64decode(f["content_base64"]) generated_files.append( { @@ -830,8 +836,8 @@ print('No executable skill module found') def _attach_files_to_response( self, response: Any, - generated_files: list[dict[str, Any]], - ) -> Any: + generated_files: list[dict[str, object]], + ) -> LLMResponseTypes: """ Attach generated files to the response object. @@ -841,11 +847,13 @@ print('No executable skill module found') if not generated_files: return response + raw_response: Final = response + # Handle dict response (Anthropic/messages API format) if isinstance(response, dict): response["_litellm_generated_files"] = generated_files verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to dict response", len(generated_files)) - return response + return raw_response # Handle object response (OpenAI format) try: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f2cb1124fa0..2de1d177b33 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,7 +18,7 @@ import os import re import secrets import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast @@ -171,8 +171,12 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ... + async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ... + async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ... + async def update( self, *, @@ -181,6 +185,10 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): ) -> _PrismaRowT | None: ... +class _TxTables(Protocol): + litellm_proxymodeltable: _PrismaTableActions[object] + + def _prisma_table( repository: BaseRepository[_RepositoryModelT], ) -> _PrismaTableActions[_RepositoryModelT]: @@ -1650,9 +1658,12 @@ async def generate_key_fn( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - if user_custom_key_generate is not None: - if inspect.iscoroutinefunction(user_custom_key_generate): - result: Final = await user_custom_key_generate(data) + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( + user_custom_key_generate + ) + if custom_key_generate_hook is not None: + if inspect.iscoroutinefunction(custom_key_generate_hook): + result: Final = await custom_key_generate_hook(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision: Final = result.get("decision", True) @@ -1847,9 +1858,10 @@ async def generate_service_account_key_fn( verbose_proxy_logger.debug("entered /key/generate") - if user_custom_key_generate is not None: - if inspect.iscoroutinefunction(user_custom_key_generate): - result: Final = await user_custom_key_generate(data) + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate + if custom_key_generate_hook is not None: + if inspect.iscoroutinefunction(custom_key_generate_hook): + result: Final = await custom_key_generate_hook(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision: Final = result.get("decision", True) @@ -1918,7 +1930,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) casted_metadata[reserved_field] = existing_value - data_json: Final = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True) try: for k, v in data_json.items(): @@ -2179,7 +2191,7 @@ async def _process_single_key_update( llm_router: Router | None, user_custom_key_update: Callable | None = None, existing_key_row: LiteLLM_VerificationToken | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2722,9 +2734,10 @@ async def update_key_fn( ) # Custom key update hook - if user_custom_key_update is not None: - if inspect.iscoroutinefunction(user_custom_key_update): - result: Final = await user_custom_key_update(data) + custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update + if custom_key_update_hook is not None: + if inspect.iscoroutinefunction(custom_key_update_hook): + result: Final = await custom_key_update_hook(data) else: raise ValueError("user_custom_key_update must be a coroutine") decision: Final = result.get("decision", True) @@ -4089,10 +4102,11 @@ async def delete_verification_tokens( failed_tokens: list = [] try: if prisma_client: - tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"token": {"in": tokens}}) + hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens] + tokens = hashed_tokens + _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await _prisma_table( + VerificationTokenRepository(prisma_client) + ).find_many(where={"token": {"in": hashed_tokens}}) if len(_keys_being_deleted) == 0: raise HTTPException( @@ -4291,7 +4305,7 @@ async def _rotate_master_key( if models: decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models) verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models)) - new_models: Final = [] + new_models: Final[list[dict[str, object]]] = [] for model in decrypted_models: new_model = await _add_model_to_db( model_params=Deployment(**model), @@ -4306,7 +4320,8 @@ async def _rotate_master_key( _dumped["model_info"] = prisma.Json(_dumped["model_info"]) new_models.append(_dumped) verbose_proxy_logger.debug("Resetting proxy model table") - async with prisma_client.db.tx() as tx: + async with prisma_client.db.tx() as tx_ctx: + tx: Final[_TxTables] = tx_ctx await tx.litellm_proxymodeltable.delete_many() verbose_proxy_logger.debug("Creating %s models", len(new_models)) await tx.litellm_proxymodeltable.create_many( @@ -4630,7 +4645,7 @@ async def _execute_virtual_key_regeneration( _validate_key_alias_format(key_alias=new_key_alias) verbose_proxy_logger.debug("non_default_values: %s", non_default_values) update_data.update(non_default_values) - update_data = prisma_client.jsonify_object(data=update_data) + jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) # If grace period set, insert deprecated key so old key remains valid await _insert_deprecated_key( @@ -4642,9 +4657,9 @@ async def _execute_virtual_key_regeneration( updated_token: Final = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, - data=update_data, + data=jsonified_update_data, ) - updated_token_dict: Final = dict(updated_token) if updated_token is not None else {} + updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {} updated_token_dict["key"] = new_token updated_token_dict["token_id"] = updated_token_dict.pop("token") @@ -5589,7 +5604,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -5602,7 +5617,7 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 @@ -5695,7 +5710,7 @@ def _build_key_filter_conditions( agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, -) -> dict[str, str | dict[str, Any] | list[dict[str, Any]]]: +) -> Mapping[str, object]: """Build filter conditions for key listing. Visibility rules: @@ -5707,14 +5722,14 @@ def _build_key_filter_conditions( so former members cannot see service accounts they created after leaving. """ # Prepare filter conditions - where: dict[str, str | dict[str, Any] | list[dict[str, Any]]] = {} + where: dict[str, object] = {} where.update(_get_condition_to_filter_out_ui_session_tokens()) # Build the OR conditions for user's keys and admin team keys - or_conditions: Final[list[dict[str, Any]]] = [] + or_conditions: Final[list[dict[str, object]]] = [] # Base conditions for user's own keys - user_condition: Final[dict[str, Any]] = {} + user_condition: Final[dict[str, object]] = {} if user_id and isinstance(user_id, str): if use_substring_matching: user_condition["user_id"] = { @@ -5784,7 +5799,7 @@ def _build_key_filter_conditions( # Apply team_id, project_id and access_group_id as global AND filters so they # narrow results across all visibility conditions (own keys, team keys, etc.) - global_filters: tuple[dict[str, Any], ...] = ( + global_filters: Final[tuple[dict[str, object], ...]] = ( *( ( {"key_alias": {"contains": key_alias, "mode": "insensitive"}} @@ -5805,7 +5820,7 @@ def _build_key_filter_conditions( else () ), ) - combined_where = {"AND": [where, *global_filters]} if global_filters else where + combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where verbose_proxy_logger.debug("Filter conditions: %s", combined_where) return combined_where @@ -5986,7 +6001,7 @@ async def _list_key_helper( ) -def _get_condition_to_filter_out_ui_session_tokens() -> dict[str, Any]: +def _get_condition_to_filter_out_ui_session_tokens() -> Mapping[str, object]: """ Condition to filter out UI session tokens """ @@ -6395,7 +6410,7 @@ async def _can_user_query_key_info( async def test_key_logging( user_api_key_dict: UserAPIKeyAuth, request: Request, - key_logging: list[dict[str, Any]], + key_logging: Sequence[Mapping[str, str]], ) -> LoggingCallbackStatus: """ Test the key-based logging diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 71407c89813..c2087005863 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,9 +13,9 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from json import JSONDecodeError -from typing import Any, Final, Literal, cast +from typing import Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError @@ -78,9 +78,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, - DeploymentTypedDict, GenericLiteLLMParams, - LiteLLMParamsTypedDict, updateDeployment, ) from litellm.utils import get_utc_datetime @@ -104,10 +102,80 @@ class UpdatePublicModelGroupsRequest(BaseModel): model_config = ConfigDict(extra="forbid") +class _ProxyModelRow(Protocol): + model_id: str + model_name: str + model_info: Mapping[str, object] | None + + def model_dump_json(self, *, exclude_none: bool = False) -> str: ... + + +class _ProxyModelTable(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ... + + def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ... + + def delete(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ... + + def delete_many(self, *, where: Mapping[str, object]) -> Awaitable[int]: ... + + +class _TxModelTables(Protocol): + litellm_proxymodeltable: _ProxyModelTable + + +class _TeamRow(Protocol): + models: Sequence[str] + + def model_dump(self) -> Mapping[str, object]: ... + + +class _TeamTable(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ... + + def update( + self, *, where: Mapping[str, object], data: Mapping[str, object], include: Mapping[str, bool] + ) -> Awaitable[LiteLLM_TeamTable]: ... + + +class _TeamIdRef(Protocol): + team_id: str + + +class _ModelAliasRow(Protocol): + id: int + model_aliases: dict[str, str] + team: _TeamIdRef | None + + +class _ModelAliasTable(Protocol): + def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRow]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ... + + +def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable: + return ModelRepository(prisma_client).table + + +def _repo_team_table(prisma_client: PrismaClient) -> _TeamTable: + return TeamRepository(prisma_client).table + + +def _db_team_table(prisma_client: PrismaClient) -> _TeamTable: + return prisma_client.db.litellm_teamtable + + +def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable: + return ModelTableRepository(prisma_client).table + + async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None: db_model: Final = cast( BaseModel | None, - await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}), + await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id}), ) if not db_model: @@ -166,14 +234,9 @@ def _raise_on_strategy_router_write_violation( def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: - merged_deployment_dict: Final = DeploymentTypedDict( - model_name=db_model.model_name, - litellm_params=LiteLLMParamsTypedDict(**db_model.litellm_params.model_dump(exclude_none=True)), - model_info=db_model.model_info.model_dump(exclude_none=True), - ) - # update model name - if updated_patch.model_name: - merged_deployment_dict["model_name"] = updated_patch.model_name + merged_model_name: Final = updated_patch.model_name or db_model.model_name + merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True) + merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True) # update litellm params if updated_patch.litellm_params: @@ -182,13 +245,11 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() } - merged_deployment_dict["litellm_params"].update(encrypted_params) + merged_litellm_params.update(encrypted_params) # update model info if updated_patch.model_info: - if "model_info" not in merged_deployment_dict: - merged_deployment_dict["model_info"] = {} - merged_deployment_dict["model_info"].update(updated_patch.model_info.model_dump(exclude_none=True)) + merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI # passes through (which today re-sends the OLD pricing on every save) cannot @@ -202,29 +263,25 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.litellm_params: for field in updated_patch.litellm_params.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: - merged_deployment_dict["litellm_params"].pop(field, None) - merged_deployment_dict.get("model_info", {}).pop(field, None) + merged_litellm_params.pop(field, None) + merged_model_info.pop(field, None) if updated_patch.model_info: for field in updated_patch.model_info.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: - merged_deployment_dict["model_info"].pop(field, None) - merged_deployment_dict.get("litellm_params", {}).pop(field, None) + merged_model_info.pop(field, None) + merged_litellm_params.pop(field, None) # convert to prisma compatible format - prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel() - if "model_name" in merged_deployment_dict: - prisma_compatible_model_dict["model_name"] = merged_deployment_dict["model_name"] + for key, value in merged_model_info.items(): + if isinstance(value, datetime.datetime): + merged_model_info[key] = value.isoformat() - if "litellm_params" in merged_deployment_dict: - prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"]) - - if "model_info" in merged_deployment_dict: - model_info: Final = merged_deployment_dict["model_info"] - for key, value in model_info.items(): - if isinstance(value, datetime.datetime): - model_info[key] = value.isoformat() - prisma_compatible_model_dict["model_info"] = json.dumps(model_info) + prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( + model_name=merged_model_name, + litellm_params=json.dumps(merged_litellm_params), + model_info=json.dumps(merged_model_info), + ) if updated_patch.blocked is not None: prisma_compatible_model_dict["blocked"] = updated_patch.blocked @@ -338,7 +395,7 @@ async def patch_model( update_data["updated_at"] = cast(str, get_utc_datetime()) # Perform partial update - updated_model: Final = await ModelRepository(prisma_client).table.update( + updated_model: Final = await _proxy_model_table(prisma_client).update( where={"model_id": model_id}, data=update_data, ) @@ -769,8 +826,8 @@ async def _setup_new_team_model_assignment( async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient, table: Any | None = None -) -> list[LiteLLM_ProxyModelTable]: + team_id: str, prisma_client: PrismaClient, table: _ProxyModelTable | None = None +) -> Sequence[_ProxyModelRow]: """ Fetch all deployments for a given team_id from the database. @@ -785,7 +842,7 @@ async def _get_team_deployments( existing transaction. """ prefix: Final = f"model_name_{team_id}_" - table = table or ModelRepository(prisma_client).table + table = table or _proxy_model_table(prisma_client) response: Final = await table.find_many( where={ "model_name": {"startswith": prefix}, @@ -806,7 +863,7 @@ async def _get_team_deployments( async def delete_team_models( team_ids: list[str], prisma_client: PrismaClient, - llm_router: Any | None, + llm_router: Router | None, ) -> list[str]: """ Delete every BYOK model owned by the given teams, from the DB and the router. @@ -820,7 +877,8 @@ async def delete_team_models( Returns the model_ids that were deleted. """ deleted_model_ids: Final[list[str]] = [] - async with prisma_client.db.tx() as tx: + async with prisma_client.db.tx() as tx_ctx: + tx: Final[_TxModelTables] = tx_ctx for team_id in team_ids: rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] @@ -920,11 +978,11 @@ async def _remove_unbacked_team_models( if not names_to_remove: return - existing_team_row: Final = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + existing_team_row: Final = await _db_team_table(prisma_client).find_unique(where={"team_id": team_id}) if existing_team_row is None: return - updated_team_row: Final[LiteLLM_TeamTable] = await prisma_client.db.litellm_teamtable.update( + updated_team_row: Final[LiteLLM_TeamTable] = await _db_team_table(prisma_client).update( where={"team_id": team_id}, data={"models": [model for model in existing_team_row.models if model not in names_to_remove]}, include={"object_permission": True}, @@ -953,7 +1011,7 @@ async def _update_existing_team_model_assignment( """ def _get_team_public_model_name( - model_info: dict | str | None, + model_info: object, ) -> str | None: parsed: Final = model_info_as_mapping(model_info) if parsed is None: @@ -1062,7 +1120,7 @@ class ModelManagementAuthChecks: detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row: Final = await TeamRepository(prisma_client).table.find_unique( + _existing_team_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1091,7 +1149,7 @@ class ModelManagementAuthChecks: ) -> Literal[True]: ## Check team model auth if model_params.model_info is not None and model_params.model_info.team_id is not None: - team_obj_row: Final = await TeamRepository(prisma_client).table.find_unique( + team_obj_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: @@ -1192,7 +1250,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result: Final = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id}) + result: Final = await _proxy_model_table(prisma_client).delete(where={"model_id": model_info.id}) if result is None: raise HTTPException( @@ -1265,9 +1323,9 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases: Final = await ModelTableRepository(prisma_client).table.find_many(include={"team": True}) + team_model_aliases: Final = await _model_alias_table(prisma_client).find_many(include={"team": True}) tasks: Final = [] - removed_model_aliases: Final = [] + removed_model_aliases: Final[list[tuple[str, str]]] = [] for team_model_alias in team_model_aliases: model_aliases = team_model_alias.model_aliases # {"alias": "public model name"} id = team_model_alias.id @@ -1278,7 +1336,7 @@ async def delete_team_model_alias( removed_model_aliases.append((team_model_alias.team.team_id, key)) del model_aliases[key] tasks.append( - ModelTableRepository(prisma_client).table.update( + _model_alias_table(prisma_client).update( where={"id": id}, data={"model_aliases": json.dumps(model_aliases)}, ) @@ -1492,7 +1550,7 @@ async def update_model( }, ) - _model_id = None + _model_id: str | None = None _model_info: Final = getattr(model_params, "model_info", None) if _model_info is None: raise Exception("model_info not provided") @@ -1551,11 +1609,11 @@ async def update_model( else: pass - _data: Final[dict] = { + _data: Final[dict[str, str]] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response: Final = await ModelRepository(prisma_client).table.update( + model_response: Final = await _proxy_model_table(prisma_client).update( where={"model_id": _model_id}, data=_data, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0b99879f9fe..97f494c51de 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -15,11 +15,11 @@ import math import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Annotated, Final, Protocol, TypeVar, cast +from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import litellm from litellm._logging import verbose_proxy_logger @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( UI_TEAM_ID, BlockTeamRequest, + BudgetNewRequest, CommonProxyErrors, DeleteTeamRequest, LiteLLM_AccessGroupTable, @@ -156,6 +157,15 @@ router: Final = APIRouter() _DbRecordT = TypeVar("_DbRecordT") +class _TeamIdKeyCount(TypedDict): + team_id: int + + +class _TeamIdGroupRow(TypedDict): + team_id: str + _count: _TeamIdKeyCount + + class _PrismaTableActions(Protocol[_DbRecordT]): async def find_unique( self, @@ -220,59 +230,127 @@ class _PrismaTableActions(Protocol[_DbRecordT]): where: Mapping[str, object] | None = None, ) -> int: ... + async def group_by( + self, + by: Sequence[str], + where: Mapping[str, object] | None = None, + count: Mapping[str, bool] | None = None, + ) -> Sequence[_TeamIdGroupRow]: ... + + +class _HasTableActions(Protocol[_DbRecordT]): + @property + def table(self) -> "_PrismaTableActions[_DbRecordT]": ... + + +def _typed_table( + repo: "_HasTableActions[_DbRecordT]", record_type: type[_DbRecordT] +) -> "_PrismaTableActions[_DbRecordT]": + return repo.table + + +def _as_object(value: object) -> object: + return value + + +def _nullable(value: _DbRecordT | None) -> _DbRecordT | None: + return value + + +class _UserIdRow(Protocol): + @property + def user_id(self) -> str | None: ... + + +class _HasUserIdTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_UserIdRow]": ... + + +def _user_id_rows_db(repo: "_HasUserIdTable") -> "_PrismaTableActions[_UserIdRow]": + return repo.table + + +class _RawTeamRow(Protocol): + @property + def members_with_roles(self) -> Sequence[Mapping[str, object]] | None: ... + + +class _HasRawTeamTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_RawTeamRow]": ... + + +def _raw_team_db(repo: "_HasRawTeamTable") -> "_PrismaTableActions[_RawTeamRow]": + return repo.table + + +class _BudgetWriteCall(Protocol): + async def __call__( + self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth + ) -> LiteLLM_BudgetTableFull: ... + + +def _as_budget_write(fn: "_BudgetWriteCall") -> "_BudgetWriteCall": + return fn + + +class _TeamFindManyArgs(TypedDict, total=False): + take: int + skip: int + order: Mapping[str, str] + cursor: Mapping[str, object] + + +class _TeamUiViewFilters(TypedDict, total=False): + team_id: Mapping[str, str] + team_alias: Mapping[str, str] + + +class _TeamIdInFilter(TypedDict, total=False): + team_id: Mapping[str, Sequence[str]] + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": - team_table: Final[_PrismaTableActions[LiteLLM_TeamTable]] = TeamRepository(prisma_client).table - return team_table + return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]": - membership_table: Final[_PrismaTableActions[LiteLLM_TeamMembership]] = TeamMembershipRepository(prisma_client).table - return membership_table + return _typed_table(TeamMembershipRepository(prisma_client), LiteLLM_TeamMembership) def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]": - user_table: Final[_PrismaTableActions[LiteLLM_UserTable]] = UserRepository(prisma_client).table - return user_table + return _typed_table(UserRepository(prisma_client), LiteLLM_UserTable) def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]": - model_table: Final[_PrismaTableActions[LiteLLM_ModelTable]] = ModelTableRepository(prisma_client).table - return model_table + return _typed_table(ModelTableRepository(prisma_client), LiteLLM_ModelTable) def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]": - org_table: Final[_PrismaTableActions[LiteLLM_OrganizationTable]] = OrganizationRepository(prisma_client).table - return org_table + return _typed_table(OrganizationRepository(prisma_client), LiteLLM_OrganizationTable) def _org_membership_db( prisma_client: PrismaClient | None, ) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]": - org_membership_table: _PrismaTableActions[LiteLLM_OrganizationMembershipTable] = OrganizationMembershipRepository( - prisma_client - ).table - return org_membership_table + return _typed_table(OrganizationMembershipRepository(prisma_client), LiteLLM_OrganizationMembershipTable) def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]": - budget_table: Final[_PrismaTableActions[LiteLLM_BudgetTableFull]] = BudgetRepository(prisma_client).table - return budget_table + return _typed_table(BudgetRepository(prisma_client), LiteLLM_BudgetTableFull) def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]": - deleted_team_table: _PrismaTableActions[LiteLLM_DeletedTeamTable] = DeletedTeamRepository(prisma_client).table - return deleted_team_table + return _typed_table(DeletedTeamRepository(prisma_client), LiteLLM_DeletedTeamTable) def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]": - access_group_table: _PrismaTableActions[LiteLLM_AccessGroupTable] = AccessGroupRepository(prisma_client).table - return access_group_table + return _typed_table(AccessGroupRepository(prisma_client), LiteLLM_AccessGroupTable) def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]": - tokens_table: _PrismaTableActions[LiteLLM_VerificationToken] = VerificationTokenRepository(prisma_client).table - return tokens_table + return _typed_table(VerificationTokenRepository(prisma_client), LiteLLM_VerificationToken) def _sanitize_for_log(value: object) -> str: @@ -408,7 +486,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - team_member_budget_table: Final = await new_budget( + team_member_budget_table: Final = await _as_budget_write(new_budget)( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -456,7 +534,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - budget_row: Final = await update_budget( + budget_row: Final = await _as_budget_write(update_budget)( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -571,7 +649,7 @@ class TeamMemberBudgetHandler: ) if missing: - await TeamMembershipRepository(prisma_client).table.create_many( + await _team_membership_db(prisma_client).create_many( data=missing, skip_duplicates=True, # safety net against concurrent races ) @@ -1407,9 +1485,10 @@ async def new_team( complete_team_data_dict["metadata"] = encrypt_callback_vars(complete_team_data_dict["metadata"]) complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) + team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict - team_row: Final[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.create( - data=complete_team_data_dict, + team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create( + data=team_creation_data, include={"litellm_model_table": True}, ) @@ -1856,7 +1935,7 @@ async def update_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team_row is None: raise HTTPException( @@ -1884,7 +1963,7 @@ async def update_team( ) if data.max_budget is not None: - existing_soft_budget: Final = getattr(existing_team_row, "soft_budget", None) + existing_soft_budget: Final[object] = _as_object(getattr(existing_team_row, "soft_budget", None)) soft_budget_to_check: Final = data.soft_budget if data.soft_budget is not None else existing_soft_budget if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)): if data.max_budget <= soft_budget_to_check: @@ -1943,7 +2022,7 @@ async def update_team( data.organization_id = None # check org team limits - if updating team that belongs to an org - org_id_to_check: Final = ( + org_id_to_check: Final[object] = _as_object( data.organization_id if data.organization_id is not None else existing_team_row.organization_id ) if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None: @@ -1976,7 +2055,7 @@ async def update_team( TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"]) if "metadata" in updated_kv: - stored_metadata: Final = ( + stored_metadata: Final[Mapping[str, JsonValue] | None] = ( { # mutable-ok: the validator payload's isinstance guard requires a plain dict key: value for key, value in existing_team_row.metadata.items() @@ -2079,16 +2158,19 @@ async def update_team( updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Final[LiteLLM_TeamTable | None] = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, - data=updated_kv, - # `object_permission` is included so `_refresh_cached_team` - # doesn't write a cached team with the relation nulled out — - # see team_model_add for the full rationale. - include={ - "litellm_model_table": True, - "object_permission": True, - }, + team_update_data: Final[Mapping[str, object]] = updated_kv + team_row: Final[LiteLLM_TeamTable | None] = _nullable( + await _team_db(prisma_client).update( + where={"team_id": data.team_id}, + data=team_update_data, + # `object_permission` is included so `_refresh_cached_team` + # doesn't write a cached team with the relation nulled out — + # see team_model_add for the full rationale. + include={ + "litellm_model_table": True, + "object_permission": True, + }, + ) ) if team_row is None or team_row.team_id is None: @@ -2603,7 +2685,7 @@ async def _resolve_existing_member_user_ids( if not requested_user_ids: return frozenset() - found: Final = await UserRepository(prisma_client).table.find_many( + found: Final = await _user_id_rows_db(UserRepository(prisma_client)).find_many( where={ # mutable-ok: Prisma query filters are dict-shaped "user_id": { # mutable-ok: Prisma query filters are dict-shaped "in": sorted(requested_user_ids) @@ -3098,7 +3180,9 @@ async def team_member_delete( key_val["user_id"] = data.user_id elif data.user_email is not None: key_val["user_email"] = data.user_email - existing_user_rows: Final = await UserRepository(prisma_client).table.find_many(where=key_val) + existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many( + where=key_val + ) if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0): for existing_user in existing_user_rows: @@ -3106,7 +3190,7 @@ async def team_member_delete( if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - await UserRepository(prisma_client).table.update( + await _user_db(prisma_client).update( where={ "user_id": existing_user.user_id, }, @@ -3114,7 +3198,7 @@ async def team_member_delete( ) # Also clean up any existing team membership rows for this user and team - user_ids_to_delete: Final = set() + user_ids_to_delete: Final = set[str]() if data.user_id is not None: user_ids_to_delete.add(data.user_id) if existing_user_rows is not None and isinstance(existing_user_rows, list): @@ -3123,9 +3207,7 @@ async def team_member_delete( user_ids_to_delete.add(existing_user.user_id) for _uid in user_ids_to_delete: - await TeamMembershipRepository(prisma_client).table.delete_many( - where={"team_id": data.team_id, "user_id": _uid} - ) + await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid}) ## DELETE KEYS CREATED BY USER FOR THIS TEAM if user_ids_to_delete: @@ -3134,9 +3216,7 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many( + keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3151,7 +3231,7 @@ async def team_member_delete( litellm_changed_by=None, ) - await VerificationTokenRepository(prisma_client).table.delete_many( + await _tokens_db(prisma_client).delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3311,7 +3391,7 @@ async def team_member_update( ### upsert new budget budget_patch: Final = _build_member_budget_patch(data) - async with prisma_client.db.tx() as tx: + async with prisma_client.tx() as tx: await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, @@ -3654,7 +3734,7 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + keys_to_delete: list[LiteLLM_VerificationToken] = await _tokens_db(prisma_client).find_many( where={"team_id": {"in": data.team_ids}} ) @@ -4469,7 +4549,7 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped: Final = await VerificationTokenRepository(prisma_client).table.group_by( + grouped: Final = await _tokens_db(prisma_client).group_by( by=["team_id"], where={"team_id": {"in": page_team_ids}}, count={"team_id": True}, @@ -4786,7 +4866,7 @@ async def _authorize_and_filter_teams( if allowed_org_ids is not None: # Org admin: query DB for teams in their orgs - org_teams: Final = await TeamRepository(prisma_client).table.find_many( + org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4800,7 +4880,9 @@ async def _authorize_and_filter_teams( ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response: Final = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}) + response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( + include={"litellm_model_table": True} + ) return [ team for team in response @@ -4808,7 +4890,7 @@ async def _authorize_and_filter_teams( ] else: # Proxy admin: all teams - return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})) + return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True})) @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -4860,7 +4942,7 @@ async def list_team( _team_memberships.append(tm) # add all keys that belong to the team - keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id}) + keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id}) try: returned_responses.append( @@ -4911,7 +4993,7 @@ async def get_paginated_teams( total_count: Final = await _team_db(prisma_client).count() # Get paginated teams - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( skip=skip, take=page_size, order={"team_alias": "asc"}, # Sort by team_alias @@ -4961,7 +5043,7 @@ async def ui_view_teams( skip: Final = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Final = {} + where_conditions: Final[_TeamUiViewFilters] = {} if team_id: where_conditions["team_id"] = { @@ -4976,7 +5058,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -5166,13 +5248,13 @@ async def team_model_delete( ) # Get current models list - current_models: Final = team_obj.models or [] + current_models: Final[Sequence[str]] = team_obj.models or [] # Remove specified models updated_models: Final = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team: Final = await TeamRepository(prisma_client).table.update( + updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, @@ -5425,7 +5507,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi BATCH_SIZE: Final = 500 while True: - find_args: dict = { + find_args: _TeamFindManyArgs = { "take": BATCH_SIZE, "order": {"team_id": "asc"}, } @@ -5433,7 +5515,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi find_args["cursor"] = {"team_id": cursor} find_args["skip"] = 1 - teams = await TeamRepository(prisma_client).table.find_many(**find_args) + teams = await _team_db(prisma_client).find_many(**find_args) if not teams: break @@ -5528,11 +5610,11 @@ async def get_team_daily_activity( ) ## Fetch team aliases and check team admin status - where_condition: Final = {} + where_condition: Final[_TeamIdInFilter] = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases: Final = await TeamRepository(prisma_client).table.find_many(where=where_condition) - team_alias_metadata: Final = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} + team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition) + team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases} # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys. diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 44abc56713f..a2c50590dd5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,9 +16,22 @@ import json import os import re import secrets +from collections.abc import Mapping, Sequence from copy import deepcopy from html import escape -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NoReturn, + Optional, + Protocol, + TypeVar, + Union, + cast, + overload, +) from urllib.parse import parse_qs, urlencode, urlparse if TYPE_CHECKING: @@ -155,6 +168,102 @@ _CLI_SSO_SECRET_KEY_FRAGMENTS: Final = frozenset( } ) +_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True) + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique( + self, + where: Mapping[str, object], + ) -> _DbRecordT | None: ... + + async def find_first( + self, + where: Mapping[str, object] | None = None, + ) -> _DbRecordT | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + ) -> Sequence[_DbRecordT]: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _DbRecordT: ... + + async def update_many( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> int: ... + + +class _UserMetadataRow(Protocol): + @property + def metadata(self) -> Mapping[str, object] | None: ... + + +class _HasUserMetadataTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_UserMetadataRow]": ... + + +def _user_meta_db(repo: "_HasUserMetadataTable") -> "_PrismaTableActions[_UserMetadataRow]": + return repo.table + + +class _SsoConfigRow(Protocol): + @property + def sso_settings(self) -> Mapping[str, object] | None: ... + + +class _HasSsoConfigTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_SsoConfigRow]": ... + + +def _sso_config_db(repo: "_HasSsoConfigTable") -> "_PrismaTableActions[_SsoConfigRow]": + return repo.table + + +class _TeamDetailRow(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _HasTeamDetailTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_TeamDetailRow]": ... + + +def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDetailRow]": + return repo.table + + +class _CustomSsoCall(Protocol): + async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... + + +class _ServicePrincipalAssignment(Protocol): + def get(self, key: str) -> str: ... + + +class _ServicePrincipalPage(Protocol): + @overload + def get( + self, + key: Literal["value"], + default: Sequence["_ServicePrincipalAssignment"], + ) -> Sequence["_ServicePrincipalAssignment"]: ... + + @overload + def get(self, key: Literal["@odata.nextLink"]) -> str | None: ... + + +def _as_object(value: object) -> object: + return value + def _hash_cli_sso_secret(secret: str) -> str: return hashlib.sha256(secret.encode("utf-8")).hexdigest() @@ -256,7 +365,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: flow = cache.get_cache(key=cache_key) if isinstance(flow, str): try: - flow = json.loads(flow) + flow = _as_object(json.loads(flow)) except ValueError: flow = None if not isinstance(flow, dict) or "poll_secret_hash" not in flow: @@ -421,7 +530,7 @@ def _flatten_cli_sso_metadata_for_poll( def build_cli_sso_attribution_metadata( result: CustomOpenID | OpenID | dict, -) -> dict[str, Any]: +) -> dict[str, object]: """ Build allowlisted, non-secret scalar attribution metadata from an SSO result. @@ -432,7 +541,7 @@ def build_cli_sso_attribution_metadata( if not claim_map: return {} - metadata: Final[dict[str, Any]] = {} + metadata: Final[dict[str, object]] = {} for source_claim, dest_key in claim_map: if not _is_safe_cli_sso_metadata_dest_key(dest_key): verbose_proxy_logger.debug("Skipping unsafe CLI SSO metadata destination key: %s", dest_key) @@ -474,14 +583,14 @@ def _merge_cli_sso_attribution_metadata( async def _persist_cli_sso_user_metadata( prisma_client: PrismaClient, user_id: str, - attribution_metadata: dict[str, Any], + attribution_metadata: dict[str, object], ) -> None: if not attribution_metadata: return try: - user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - existing_metadata: dict[str, Any] = {} + user_row: Final = await _user_meta_db(UserRepository(prisma_client)).find_unique(where={"user_id": user_id}) + existing_metadata: dict[str, object] = {} if user_row is not None: row_metadata: Final = user_row.metadata if isinstance(row_metadata, dict): @@ -491,7 +600,7 @@ async def _persist_cli_sso_user_metadata( existing_metadata=existing_metadata, attribution_metadata=attribution_metadata, ) - await UserRepository(prisma_client).table.update_many( + await _user_meta_db(UserRepository(prisma_client)).update_many( where={"user_id": user_id}, data={"metadata": merged_metadata}, ) @@ -1104,7 +1213,7 @@ def generic_response_convertor( ) # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified - extra_fields: dict[str, Any] | None = None + extra_fields: dict[str, object] | None = None if generic_user_extra_attributes: extra_fields = {} for attr_name in generic_user_extra_attributes.split(","): @@ -1193,7 +1302,9 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]: prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict: Final = dict(sso_db_record.sso_settings) @@ -1225,7 +1336,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict: Final = dict(sso_db_record.sso_settings) @@ -1273,7 +1386,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: return role_mappings -def _parse_generic_sso_headers() -> dict: +def _parse_generic_sso_headers() -> dict[str, str]: """Parse comma-separated GENERIC_SSO_HEADERS env var into a dict.""" raw: Final = os.getenv("GENERIC_SSO_HEADERS", None) if raw is None: @@ -1677,7 +1790,7 @@ def _build_sso_user_update_data( result: Union["CustomOpenID", OpenID, dict] | None, user_email: str | None, user_id: str | None, -) -> dict: +) -> dict[str, object]: """ Build the update data dictionary for SSO user upsert. @@ -1689,7 +1802,7 @@ def _build_sso_user_update_data( Returns: dict: Update data containing user_email and optionally user_role if valid """ - update_data: Final[dict] = {"user_email": normalize_email(user_email)} + update_data: Final[dict[str, object]] = {"user_email": normalize_email(user_email)} # Get SSO role from result and include if valid sso_role: Final = getattr(result, "user_role", None) @@ -1740,7 +1853,7 @@ async def _sync_user_role_from_jwt_role_map( # Update existing DB record if role differs if user_info is not None and user_info.user_role != mapped_role.value: - await UserRepository(prisma_client).table.update( + await _user_meta_db(UserRepository(prisma_client)).update( where={"user_id": user_info.user_id}, data={"user_role": mapped_role.value}, ) @@ -1796,7 +1909,7 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism return user_role if prisma_client: - await UserRepository(prisma_client).table.update( + await _user_meta_db(UserRepository(prisma_client)).update( where={"user_id": user_id}, data={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, ) @@ -2016,10 +2129,11 @@ async def _build_cli_sso_user_defined_values( ) -> SSOUserDefinedValues | None: from litellm.proxy.proxy_server import user_custom_sso + custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso user_id: Final = parsed_openid_result.get("user_id") - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - return await user_custom_sso(result) + if custom_sso_handler is not None: + if inspect.iscoroutinefunction(custom_sso_handler): + return await custom_sso_handler(result) raise ValueError("user_custom_sso must be a coroutine function") if user_id is None: return None @@ -2035,12 +2149,14 @@ async def _build_cli_sso_user_defined_values( async def _fetch_cli_sso_team_details( prisma_client: PrismaClient, - teams: list[str], -) -> list[dict[str, Any]]: - team_details: Final[list[dict[str, Any]]] = [] + teams: Sequence[str], +) -> list[dict[str, object]]: + team_details: Final[list[dict[str, object]]] = [] try: if teams: - prisma_teams: Final = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": teams}}) + prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many( + where={"team_id": {"in": teams}} + ) for team_row in prisma_teams: team_dict = team_row.model_dump() team_details.append( @@ -2257,12 +2373,12 @@ async def cli_poll_key( verbose_proxy_logger.info("Returning teams list for user %s to select from: %s", user_id, user_teams) # Best-effort construction of team_details if it wasn't # already cached for some reason. - team_details_response: list[dict[str, Any]] | None = None + team_details_response: list[dict[str, object]] | None = None if isinstance(user_team_details, list) and user_team_details: team_details_response = user_team_details elif user_teams: team_details_response = [{"team_id": t, "team_alias": None} for t in user_teams] - poll_response: dict[str, Any] = { + poll_response: dict[str, object] = { "status": "ready", "user_id": user_id, "teams": user_teams, @@ -2997,7 +3113,9 @@ class SSOAuthenticationHandler: user_id=user_id, ) - await UserRepository(prisma_client).table.update_many(where={"user_id": user_id}, data=update_data) + await _user_meta_db(UserRepository(prisma_client)).update_many( + where={"user_id": user_id}, data=update_data + ) else: verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB") # user not in DB, insert User into LiteLLM DB @@ -3089,7 +3207,9 @@ class SSOAuthenticationHandler: code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) try: - team_obj: Final = await TeamRepository(prisma_client).table.find_first(where={"team_id": litellm_team_id}) + team_obj: Final = await _team_detail_db(TeamRepository(prisma_client)).find_first( + where={"team_id": litellm_team_id} + ) verbose_proxy_logger.debug("Team object: %s", team_obj) # only create a new team if it doesn't exist @@ -3278,9 +3398,10 @@ class SSOAuthenticationHandler: # But if it is, we want their models preferences user_defined_values: SSOUserDefinedValues | None = None - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - user_defined_values = await user_custom_sso(result) + custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso + if custom_sso_handler is not None: + if inspect.iscoroutinefunction(custom_sso_handler): + user_defined_values = await custom_sso_handler(result) else: raise ValueError("user_custom_sso must be a coroutine function") elif user_id is not None: @@ -3448,7 +3569,7 @@ class SSOAuthenticationHandler: dict: Token exchange parameters """ # Prepare token exchange parameters (may add code_verifier: str later) - token_params: Final[dict[str, Any]] = {"include_client_id": generic_include_client_id} + token_params: Final[dict[str, object]] = {"include_client_id": generic_include_client_id} # Retrieve PKCE code_verifier if PKCE was used in authorization. # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip @@ -3663,7 +3784,7 @@ class SSOAuthenticationHandler: access_token string. Raises ProxyException on any validation failure. """ try: - token_response_raw: Final = response.json() + token_response_raw: Final[object] = _as_object(response.json()) except Exception as json_err: verbose_proxy_logger.error( "Failed to parse token response as JSON: %s. Body: %s", @@ -4253,7 +4374,7 @@ class MicrosoftSSOHandler: while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: response = await async_client.get(next_link, headers=headers) - response_json = response.json() + response_json: _ServicePrincipalPage = response.json() verbose_proxy_logger.debug("Response from service principal app role assigned to: %s", response_json) for _object in response_json.get("value", []): diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 382df608a0c..08bb8698cac 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -4,11 +4,11 @@ import json import os from collections import Counter from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, Protocol, TypeVar from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile -from pydantic import ConfigDict, ValidationError, create_model +from pydantic import ConfigDict, JsonValue, ValidationError, create_model from pydantic.fields import FieldInfo import litellm @@ -36,6 +36,73 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router: Final = APIRouter() +_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True) + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique(self, where: Mapping[str, object]) -> _DbRecordT | None: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ... + + async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ... + + +class _SsoSettingsMappingRow(Protocol): + @property + def sso_settings(self) -> Mapping[str, object] | None: ... + + +class _HasSsoSettingsMappingTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_SsoSettingsMappingRow]: ... + + +def _sso_settings_mapping_db(repo: _HasSsoSettingsMappingTable) -> _PrismaTableActions[_SsoSettingsMappingRow]: + return repo.table + + +class _StoredSsoSettingsRow(Protocol): + @property + def sso_settings(self) -> object: ... + + +class _HasStoredSsoSettingsTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_StoredSsoSettingsRow]: ... + + +def _stored_sso_settings_db(repo: _HasStoredSsoSettingsTable) -> _PrismaTableActions[_StoredSsoSettingsRow]: + return repo.table + + +class _UiSettingsRow(Protocol): + @property + def ui_settings(self) -> str | Mapping[str, JsonValue] | None: ... + + +class _HasUiSettingsTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_UiSettingsRow]: ... + + +def _ui_settings_db(repo: _HasUiSettingsTable) -> _PrismaTableActions[_UiSettingsRow]: + return repo.table + + +class _ConfigParamRow(Protocol): + @property + def param_value(self) -> str | Mapping[str, object] | None: ... + + +class _HasConfigParamTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_ConfigParamRow]: ... + + +def _config_param_db(repo: _HasConfigParamTable) -> _PrismaTableActions[_ConfigParamRow]: + return repo.table + + # Maps each UIThemeConfig field to the env var the UI branding path reads it # from. /update/ui_theme_settings writes both the stored ui_theme_config and # these env vars, so /get/ui_theme_settings resolves the same env vars to @@ -54,7 +121,7 @@ def _is_public_http_url(value: str | None) -> bool: return parsed.scheme in ("http", "https") and bool(parsed.netloc) -def _resolve_ui_theme_field(stored_values: Mapping[str, Any], field_name: str) -> str | None: +def _resolve_ui_theme_field(stored_values: Mapping[str, object], field_name: str) -> str | None: """Resolve one UI theme field to the value the branding path actually uses. The stored ui_theme_config wins; a field absent or blank there falls back to @@ -263,7 +330,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [ # include generics like ``Optional[int]`` / ``List[str]`` that are not # instances of ``type`` — so tightening this to ``type`` would reject # valid inputs. -_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[Any, FieldInfo]]] = {} +_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[object, FieldInfo]]] = {} # Settings OSS knows about as enterprise-gated. If a caller sends one of # these keys and no extension package has registered it, the PATCH @@ -275,7 +342,7 @@ _ENTERPRISE_ONLY_UI_SETTINGS: Final[set[str]] = {"enable_projects_ui"} _EFFECTIVE_UI_SETTINGS_CLASS: type[UISettings] | None = None -def register_extra_ui_setting(name: str, annotation: Any, field: FieldInfo) -> None: +def register_extra_ui_setting(name: str, annotation: object, field: FieldInfo) -> None: """Register an additional UI settings field contributed by an extension package. ``field`` must be a ``FieldInfo`` instance — construct it directly @@ -470,7 +537,7 @@ async def delete_allowed_ip( async def _get_settings_with_schema( settings_key: str, - settings_class: Any, + settings_class: type[BaseModel], config: dict, ) -> dict: """ @@ -842,7 +909,9 @@ async def get_sso_settings(): # Resolve the effective SSO config: the stored row wins, else the process # environment, else each field's default. Unlike the legacy read path this # does not write os.environ; a GET has no business mutating the environment. - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_settings_mapping_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) sso_db_settings: Final = dict(sso_db_record.sso_settings) if sso_db_record and sso_db_record.sso_settings else None resolved: Final = resolve_sso_config(sso_db_settings, os.environ) @@ -914,8 +983,10 @@ async def update_sso_settings( # before-snapshot has the same shape as after_value, and rely on # create_config_audit_log's secret-name redaction to mask the # *_client_secret fields before the audit row is written. - existing_sso_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) - before_sso_data: dict[str, Any] | None = None + existing_sso_record: Final = await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) + before_sso_data: dict[str, JsonValue] | None = None if existing_sso_record and existing_sso_record.sso_settings: stored = existing_sso_record.sso_settings if isinstance(stored, str): @@ -948,7 +1019,7 @@ async def update_sso_settings( encrypted_sso_data: Final = proxy_config._encrypt_env_variables(environment_variables=sso_data) # Save to dedicated SSO table - await SSOConfigRepository(prisma_client).table.upsert( + await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).upsert( where={"id": "sso_config"}, data={ "create": { @@ -974,7 +1045,7 @@ async def update_sso_settings( # Remove SSO-related env vars from config.environment_variables try: - env_var_entry: Final = await ConfigRepository(prisma_client).table.find_unique( + env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client)).find_unique( where={"param_name": "environment_variables"} ) @@ -982,7 +1053,7 @@ async def update_sso_settings( if env_var_entry is not None: if env_var_entry.param_value is not None: if isinstance(env_var_entry.param_value, str): - environment_variables = json.loads(env_var_entry.param_value) + environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value) else: environment_variables = dict(env_var_entry.param_value) else: @@ -993,7 +1064,7 @@ async def update_sso_settings( key: value for key, value in environment_variables.items() if key not in env_vars_to_remove } - await ConfigRepository(prisma_client).table.update( + await _config_param_db(ConfigRepository(prisma_client)).update( where={"param_name": "environment_variables"}, data={ "param_value": json.dumps(filtered_env_vars, default=str), @@ -1239,8 +1310,10 @@ async def get_ui_settings_cached() -> dict[str, Any]: if prisma_client is None: return {} - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) - ui_settings: dict[str, Any] = {} + db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) + ui_settings: dict[str, JsonValue] = {} if db_record and db_record.ui_settings: raw: Final = db_record.ui_settings ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw) @@ -1272,9 +1345,11 @@ async def get_ui_settings(): detail={"error": "Database not connected. Please connect a database."}, ) - ui_settings: dict[str, Any] = {} + ui_settings: Mapping[str, JsonValue] = {} - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) if db_record and db_record.ui_settings: ui_settings_json: Final = db_record.ui_settings @@ -1300,7 +1375,7 @@ async def get_ui_settings(): await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) # Build config-like object for schema helper - config: Final[dict[str, Any]] = {"litellm_settings": {"ui_settings": ui_settings}} + config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}} return await _get_settings_with_schema( settings_key="ui_settings", @@ -1315,7 +1390,7 @@ async def get_ui_settings(): dependencies=[Depends(user_api_key_auth)], ) async def update_ui_settings( - settings_body: dict[str, Any] = Body(...), + settings_body: dict[str, object] = Body(...), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1352,7 +1427,7 @@ async def update_ui_settings( raise HTTPException(status_code=422, detail=e.errors()) # Only include fields the caller actually sent (not Pydantic defaults). - settings_dict: Final = settings.model_dump(exclude_unset=True) + settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True) # Reject enterprise-only settings up front so the caller gets a clear # signal instead of a silent drop. @@ -1373,15 +1448,17 @@ async def update_ui_settings( # Merge with existing persisted settings so a partial PATCH doesn't # overwrite fields the caller didn't send. - existing: dict = {} - db_existing: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + existing: dict[str, JsonValue] = {} + db_existing: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) if db_existing and db_existing.ui_settings: raw: Final = db_existing.ui_settings existing = json.loads(raw) if isinstance(raw, str) else dict(raw) ui_settings: Final = {**existing, **incoming} - await UISettingsRepository(prisma_client).table.upsert( + await _ui_settings_db(UISettingsRepository(prisma_client)).upsert( where={"id": "ui_settings"}, data={ "create": { diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 6b40281198c..2b037bef795 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,10 +10,17 @@ All /vector_store management endpoints import copy import json -from typing import Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +if TYPE_CHECKING: + from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow + + from litellm.proxy.utils import PrismaClient + from litellm.router import Router + import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -43,6 +50,25 @@ from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router: Final = APIRouter() + +class _VectorStoreTableActions(Protocol): + async def find_unique(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ... + + async def create(self, data: Mapping[str, object]) -> "_VectorStoreRow": ... + + async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> "_VectorStoreRow": ... + + async def delete(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ... + + +def _vector_store_table(prisma_client: "PrismaClient") -> _VectorStoreTableActions: + return ManagedVectorStoresRepository(prisma_client).table + + +def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore(**row.model_dump()) + + _LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker() @@ -117,22 +143,20 @@ def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> An async def _fetch_and_authorize_vector_store( vector_store_id: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: "PrismaClient", ) -> "LiteLLM_ManagedVectorStore": """ Look up a vector store by id and confirm the caller can access it. Raises HTTPException(404) on miss and HTTPException(403) on access denial. """ - row: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( - where={"vector_store_id": vector_store_id} - ) + row: Final = await _vector_store_table(prisma_client).find_unique(where={"vector_store_id": vector_store_id}) if row is None: raise HTTPException( status_code=404, detail=f"Vector store with ID {vector_store_id} not found", ) - typed: Final = LiteLLM_ManagedVectorStore(**row.model_dump()) + typed: Final = _row_to_vector_store(row) if not await _check_vector_store_access(typed, user_api_key_dict): raise HTTPException( status_code=403, @@ -141,7 +165,7 @@ async def _fetch_and_authorize_vector_store( return typed -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, Any] | None: +def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None: """ Resolve embedding config from router's config-defined models. @@ -177,7 +201,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d litellm_params = deployment.litellm_params # Build embedding config from model params - embedding_config: dict[str, Any] = {} + embedding_config: dict[str, object] = {} # Extract api_key api_key = getattr(litellm_params, "api_key", None) @@ -217,7 +241,9 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d return None -async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) -> dict[str, Any] | None: +async def _resolve_embedding_config_from_db( + embedding_model: str, prisma_client: "PrismaClient" +) -> dict[str, object] | None: """ Resolve embedding config from database model configuration. @@ -307,7 +333,9 @@ async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) return None -async def _resolve_embedding_config(embedding_model: str, prisma_client, llm_router=None) -> dict[str, Any] | None: +async def _resolve_embedding_config( + embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None +) -> dict[str, object] | None: """ Resolve embedding config from either router (config-defined) or database models. @@ -388,7 +416,7 @@ async def _check_vector_store_access( async def create_vector_store_in_db( vector_store_id: str, custom_llm_provider: str, - prisma_client, + prisma_client: "PrismaClient | None", vector_store_name: str | None = None, vector_store_description: str | None = None, vector_store_metadata: dict | None = None, @@ -417,7 +445,7 @@ async def create_vector_store_in_db( raise HTTPException(status_code=500, detail="Database not connected") # Check if vector store already exists - existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique( where={"vector_store_id": vector_store_id} ) if existing_vector_store is not None: @@ -427,7 +455,7 @@ async def create_vector_store_in_db( ) # Prepare data for database - data_to_create: Final[dict[str, Any]] = { + data_to_create: Final[dict[str, object]] = { "vector_store_id": vector_store_id, "custom_llm_provider": custom_llm_provider, } @@ -463,9 +491,9 @@ async def create_vector_store_in_db( data_to_create["litellm_params"] = safe_dumps({}) # Create in database - _new_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.create(data=data_to_create) + _new_vector_store: Final = await _vector_store_table(prisma_client).create(data=data_to_create) - new_vector_store: Final[LiteLLM_ManagedVectorStore] = LiteLLM_ManagedVectorStore(**_new_vector_store.model_dump()) + new_vector_store: Final[LiteLLM_ManagedVectorStore] = _row_to_vector_store(_new_vector_store) # Add vector store to registry if litellm.vector_store_registry is not None: @@ -682,12 +710,12 @@ async def delete_vector_store( memory_vector_store_exists = False vector_store_to_check = None - existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique( where={"vector_store_id": data.vector_store_id} ) if existing_vector_store is not None: db_vector_store_exists = True - vector_store_to_check = LiteLLM_ManagedVectorStore(**existing_vector_store.model_dump()) + vector_store_to_check = _row_to_vector_store(existing_vector_store) # Check in-memory registry if litellm.vector_store_registry is not None: @@ -715,9 +743,7 @@ async def delete_vector_store( # Delete from database if exists if db_vector_store_exists: - await ManagedVectorStoresRepository(prisma_client).table.delete( - where={"vector_store_id": data.vector_store_id} - ) + await _vector_store_table(prisma_client).delete(where={"vector_store_id": data.vector_store_id}) # Delete from in-memory registry if exists if memory_vector_store_exists and litellm.vector_store_registry is not None: @@ -829,7 +855,7 @@ async def update_vector_store( try: update_data: Final = data.model_dump(exclude_unset=True) - vector_store_id: Final = update_data.pop("vector_store_id") + vector_store_id: Final[str] = update_data.pop("vector_store_id") # Per-store access control: anyone authenticated who passes the # premium-feature gate could otherwise update *any* vector store — @@ -857,12 +883,12 @@ async def update_vector_store( update_data["litellm_params"] = safe_dumps(litellm_params_dict) # Update in database - updated: Final = await ManagedVectorStoresRepository(prisma_client).table.update( + updated: Final = await _vector_store_table(prisma_client).update( where={"vector_store_id": vector_store_id}, data=update_data, ) - updated_vs: Final = LiteLLM_ManagedVectorStore(**updated.model_dump()) + updated_vs: Final = _row_to_vector_store(updated) # Immediately update in-memory registry to keep it in sync if litellm.vector_store_registry is not None: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f54023836e5..b2d065ea23b 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,8 +4,8 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence -from typing import Any, Final, Literal, cast +from collections.abc import Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable from openai.types.chat.chat_completion_named_tool_choice_param import ( ChatCompletionNamedToolChoiceParam, @@ -16,6 +16,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_param import FunctionToolParam +from pydantic import TypeAdapter from typing_extensions import TypedDict from litellm._logging import verbose_logger @@ -78,9 +79,35 @@ from .custom_tools import ( unwrap_custom_tool_arguments, ) +if TYPE_CHECKING: + from openai.types.responses.response_apply_patch_tool_call import ( + ResponseApplyPatchToolCall, + ) + ########### Initialize Classes used for Responses API ########### TOOL_CALLS_CACHE: Final = InMemoryCache() +_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) +_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) +_TEXT_ADAPTER: Final = TypeAdapter(str) + + +@runtime_checkable +class _SupportsIter(Protocol): + def __iter__(self) -> Iterator[object]: ... + + +@runtime_checkable +class _HasToolCalls(Protocol): + tool_calls: object + + +@runtime_checkable +class _HasId(Protocol): + id: object + class ChatCompletionSession(TypedDict, total=False): messages: list[ @@ -205,7 +232,7 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None = None, stream: bool | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, object] | None = None, **kwargs, ) -> dict: """ @@ -462,7 +489,9 @@ class LiteLLMCompletionResponsesConfig: if not chat_completion_messages: continue - deduped_in_place: list[Any] = [] + deduped_in_place: list[ + AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage + ] = [] for m in chat_completion_messages: role = "" if isinstance(m, dict): @@ -472,7 +501,7 @@ class LiteLLMCompletionResponsesConfig: # Drop assistant tool_calls wrappers if we already have this call_id if role == "assistant": - tool_calls: Any = ( + tool_calls: object = ( m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None) ) call_id = "" @@ -534,7 +563,7 @@ class LiteLLMCompletionResponsesConfig: call_id = "" if role == "assistant": - tool_calls: Any = None + tool_calls: object = None if isinstance(tool_call_message, dict): tool_calls = tool_call_message.get("tool_calls") else: @@ -578,7 +607,16 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None: + def _find_previous_assistant_idx( + messages: Sequence[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message + ], + current_idx: int, + ) -> int | None: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": @@ -586,7 +624,18 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str: + def _recover_tool_call_id_from_assistant( + assistant_message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + ) -> str: """Try to recover empty tool_call_id from assistant message's tool_calls.""" tool_calls_raw: Final = ( assistant_message.get("tool_calls") @@ -594,17 +643,23 @@ class LiteLLMCompletionResponsesConfig: else getattr(assistant_message, "tool_calls", None) ) if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0: - first_tool_call: Final = tool_calls_raw[0] + first_tool_call: Final = _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw)[0] if isinstance(first_tool_call, dict): - tool_call_id_raw = first_tool_call.get("id", "") + tool_call_id_raw = _ANY_KEY_DICT_ADAPTER.validate_python(first_tool_call).get("id", "") return str(tool_call_id_raw) if tool_call_id_raw is not None else "" - elif hasattr(first_tool_call, "id"): - tool_call_id_raw = getattr(first_tool_call, "id", None) + elif isinstance(first_tool_call, _HasId): + tool_call_id_raw = first_tool_call.id return str(tool_call_id_raw) if tool_call_id_raw is not None else "" return "" @staticmethod - def _get_tool_calls_list(assistant_message: Any) -> list[Any]: + def _get_tool_calls_list( + assistant_message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + ) -> Sequence[object]: """Extract tool_calls as a list from assistant message.""" tool_calls_raw: Final = ( assistant_message.get("tool_calls") @@ -614,18 +669,18 @@ class LiteLLMCompletionResponsesConfig: if tool_calls_raw is None: return [] if isinstance(tool_calls_raw, list): - return tool_calls_raw - if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)): + return _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw) + if isinstance(tool_calls_raw, _SupportsIter) and not isinstance(tool_calls_raw, (str, bytes)): return list(tool_calls_raw) return [] @staticmethod - def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool: + def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool: """Check if a tool_call with the given ID exists in the list.""" for tool_call in tool_calls: - tool_call_id_to_check: str | None = None + tool_call_id_to_check: object = None if isinstance(tool_call, dict): - tool_call_id_to_check = tool_call.get("id") + tool_call_id_to_check = _ANY_KEY_DICT_ADAPTER.validate_python(tool_call).get("id") elif hasattr(tool_call, "id"): tool_call_id_to_check = getattr(tool_call, "id", None) if tool_call_id_to_check == tool_call_id: @@ -633,12 +688,13 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None: + def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: Sequence[object]) -> dict[str, object] | None: """Reconstruct a minimal tool_call definition from tools list.""" for tool in tools: if isinstance(tool, dict): - tool_function = tool.get("function") or {} - tool_name = tool_function.get("name") or tool.get("name") or "" + tool_map = _ANY_KEY_DICT_ADAPTER.validate_python(tool) + tool_function = _ANY_KEY_DICT_ADAPTER.validate_python(tool_map.get("function") or {}) + tool_name = tool_function.get("name") or tool_map.get("name") or "" if tool_name: return { "id": tool_call_id, @@ -651,7 +707,7 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: + def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object: """ Safely read a field from dict-like or attribute-based objects. """ @@ -659,7 +715,7 @@ class LiteLLMCompletionResponsesConfig: return default if isinstance(obj, dict): - return obj.get(key, default) + return _ANY_KEY_DICT_ADAPTER.validate_python(obj).get(key, default) getter: Final = getattr(obj, "get", None) if callable(getter): @@ -672,13 +728,13 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _create_tool_call_chunk( - tool_use_definition: dict[str, Any], tool_call_id: str, index: int + tool_use_definition: Mapping[object, object], tool_call_id: str, index: int ) -> ChatCompletionToolCallChunk: """Create a ChatCompletionToolCallChunk from tool_use_definition.""" function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function") function_name_raw: Final = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name") function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments") - function: Final[dict[str, Any]] = { + function: Final[dict[str, object]] = { "name": function_name_raw or "", "arguments": function_arguments_raw or "{}", } @@ -697,7 +753,7 @@ class LiteLLMCompletionResponsesConfig: ) @staticmethod - def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None: + def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[object, object] | None: """ Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk. """ @@ -705,7 +761,7 @@ class LiteLLMCompletionResponsesConfig: return None if isinstance(tool_use_definition, dict): - normalized_definition: dict[str, Any] = dict(tool_use_definition) + normalized_definition: dict[object, object] = _ANY_KEY_DICT_ADAPTER.validate_python(tool_use_definition) else: tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id") tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type") @@ -738,7 +794,7 @@ class LiteLLMCompletionResponsesConfig: return normalized_definition @staticmethod - def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None: + def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None: """Add a tool_call to an assistant message.""" if isinstance(assistant_message, dict): prev_assistant_dict: Final = cast(dict[str, Any], assistant_message) @@ -747,7 +803,7 @@ class LiteLLMCompletionResponsesConfig: tool_calls_list: Final = prev_assistant_dict["tool_calls"] if isinstance(tool_calls_list, list): tool_calls_list.append(tool_call_chunk) - elif hasattr(assistant_message, "tool_calls"): + elif isinstance(assistant_message, _HasToolCalls): if assistant_message.tool_calls is None: assistant_message.tool_calls = [] if isinstance(assistant_message.tool_calls, list): @@ -762,7 +818,7 @@ class LiteLLMCompletionResponsesConfig: | ChatCompletionMessageToolCall | Message ], - tools: list[Any] | None = None, + tools: Sequence[object] | None = None, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -851,7 +907,7 @@ class LiteLLMCompletionResponsesConfig: tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list(prev_assistant) if not LiteLLMCompletionResponsesConfig._check_tool_call_exists(tool_calls, tool_call_id): - _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) + _tool_use_definition: object = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) if not _tool_use_definition and tools: _tool_use_definition = LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools( @@ -908,7 +964,7 @@ class LiteLLMCompletionResponsesConfig: function_call=input_item ) else: - content: Final = input_item.get("content") + content: Final[object] = input_item.get("content") # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages if content is None: @@ -923,7 +979,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_tool_call_output(input_item: Any) -> bool: + def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a tool call output """ @@ -936,7 +992,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_function_call(input_item: Any) -> bool: + def _is_input_item_function_call(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a function call or custom tool call. Both need to be reconstructed as assistant tool_calls for Chat @@ -946,7 +1002,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_tool_call_output_to_chat_completion_message( - tool_call_output: dict[str, Any], + tool_call_output: Mapping[str, object], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ ChatCompletionToolMessage is used to indicate the output from a tool call @@ -958,7 +1014,7 @@ class LiteLLMCompletionResponsesConfig: return [] def _normalize_function_call_output_to_tool_content( - output: Any, + output: object, ) -> Any: """ Normalize Responses API function_call_output.output into a shape that downstream @@ -981,7 +1037,7 @@ class LiteLLMCompletionResponsesConfig: # Some adapters represent tool output as a list of "input_*" parts if isinstance(output, list): - normalized_blocks: Final[list[dict[str, Any]]] = [] + normalized_blocks: Final[list[dict[str, object]]] = [] text_acc: Final[list[str]] = [] for part in output: if not isinstance(part, dict): @@ -1082,7 +1138,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_function_call_to_chat_completion_message( - function_call: dict[str, Any], + function_call: Mapping[str, str], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API function_call into a Chat Completion message with tool calls @@ -1127,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig: return [chat_completion_response_message] @staticmethod - def _resolve_file_id(item: dict[str, Any]) -> str | None: + def _resolve_file_id(item: Mapping[str, object]) -> object: """ Return the effective file_id for a Responses API input_file item. Explicit file_id takes precedence; file_url is used as fallback so @@ -1136,7 +1192,7 @@ class LiteLLMCompletionResponsesConfig: return item.get("file_id") or item.get("file_url") or None @staticmethod - def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, Any]: + def _transform_input_file_item_to_file_item(item: Mapping[str, object]) -> dict[str, object]: """ Transform a Responses API input_file item to a Chat Completion file item @@ -1146,21 +1202,21 @@ class LiteLLMCompletionResponsesConfig: Returns: Dictionary with transformed file structure for Chat Completion """ - file_dict: Final[dict[str, Any]] = {} + file_dict: Final[dict[str, object]] = {} file_id: Final = LiteLLMCompletionResponsesConfig._resolve_file_id(item) if file_id: file_dict["file_id"] = file_id if item.get("file_data"): file_dict["file_data"] = item["file_data"] - new_item: Final[dict[str, Any]] = {"type": "file", "file": file_dict} + new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: new_item["cache_control"] = item["cache_control"] return new_item @staticmethod def _transform_input_image_item_to_image_item( - item: dict[str, Any], + item: Mapping[str, str], ) -> ChatCompletionImageObject: """ Transform a Responses API input_image item to a Chat Completion image item @@ -1173,8 +1229,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_content_to_chat_completion_content( - content: Any, - ) -> str | list[str | dict[str, Any]]: + content: object, + ) -> str | list[str | dict[str, object]]: """ Transform a Responses API content into a Chat Completion content @@ -1188,7 +1244,7 @@ class LiteLLMCompletionResponsesConfig: elif isinstance(content, str): return content elif isinstance(content, list): - content_list: Final[list[str | dict[str, Any]]] = [] + content_list: Final[list[str | dict[str, object]]] = [] for item in content: if isinstance(item, str): content_list.append(item) @@ -1198,8 +1254,8 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) ) elif item.get("type") == "input_image": - image_block = dict( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) + image_block = _STR_KEY_DICT_ADAPTER.validate_python( + dict(LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item)) ) if "cache_control" in item: image_block["cache_control"] = item["cache_control"] @@ -1209,7 +1265,7 @@ class LiteLLMCompletionResponsesConfig: text_value = item.get("text") if text_value is None: continue - content_block: dict[str, Any] = { + content_block: dict[str, object] = { "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( item.get("type") or "text" ), @@ -1299,7 +1355,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(typed_tool.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - chat_completion_tool: dict[str, Any] = { + chat_completion_tool: dict[str, object] = { "type": "function", "function": { "name": typed_tool.get("name") or "", @@ -1340,7 +1396,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None, - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to Responses API request tool format. Inverse of @@ -1348,7 +1404,7 @@ class LiteLLMCompletionResponsesConfig: """ if chat_completion_tools is None or not chat_completion_tools: return [] - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for tool in chat_completion_tools: if not isinstance(tool, dict): result.append(tool) @@ -1358,7 +1414,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(fn.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - responses_tool: dict[str, Any] = { + responses_tool: dict[str, object] = { "type": "function", "name": fn.get("name") or "", "description": fn.get("description") or "", @@ -1510,7 +1566,7 @@ class LiteLLMCompletionResponsesConfig: def convert_response_function_tool_call_to_chat_completion_tool_call( tool_call_item: Any, index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseFunctionToolCall to ChatCompletionToolCallChunk format. @@ -1536,7 +1592,7 @@ class LiteLLMCompletionResponsesConfig: else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {}) ) - function_dict: Final[dict[str, Any]] = { + function_dict: Final[dict[str, object]] = { "name": tool_call_item.name, "arguments": tool_call_item.arguments, } @@ -1544,7 +1600,7 @@ class LiteLLMCompletionResponsesConfig: if provider_specific_fields: function_dict["provider_specific_fields"] = provider_specific_fields - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item( getattr(tool_call_item, "id", None), getattr(tool_call_item, "call_id", None), @@ -1561,9 +1617,9 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def convert_apply_patch_tool_call_to_chat_completion_tool_call( - tool_call_item: Any, + tool_call_item: "ResponseApplyPatchToolCall", index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. @@ -1581,7 +1637,7 @@ class LiteLLMCompletionResponsesConfig: import json operation_dict: Final = tool_call_item.operation.model_dump() - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": tool_call_item.call_id, "function": { "name": "apply_patch", @@ -1795,9 +1851,11 @@ class LiteLLMCompletionResponsesConfig: if not images: return image_generation_items - for idx, image_item in enumerate(images): + for idx, image_item in enumerate(_DICT_ITEMS_LIST_ADAPTER.validate_python(images)): # Extract base64 from data URL - image_url = image_item.get("image_url", {}).get("url", "") + image_url = _TEXT_ADAPTER.validate_python( + _ANY_KEY_DICT_ADAPTER.validate_python(image_item.get("image_url", {})).get("url", "") + ) base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) if base64_data: @@ -2048,8 +2106,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_text_format_to_response_format( - text_param: dict[str, Any] | Any, - ) -> dict[str, Any] | None: + text_param: object, + ) -> dict[str, object] | None: """ Transform Responses API text.format parameter to Chat Completion response_format parameter. diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 820839fc6bf..2e1e1a44594 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,7 +5,7 @@ import json import time import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -33,6 +33,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, + ResponseAPIUsage, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -112,7 +113,7 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]: - if isinstance(error_obj, dict): + if _is_json_object(error_obj): raw_message = error_obj.get("message") raw_type = error_obj.get("type") raw_code = error_obj.get("code") @@ -243,7 +244,9 @@ class BaseResponsesAPIStreamingIterator: # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a # truthy child Mock for any attribute, which breaks tests and is wrong on stream. if "response" in parsed_chunk: - response_object: Final = getattr(openai_responses_api_chunk, "response", None) + response_object: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_object is not None: response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=response_object, @@ -279,7 +282,9 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: - _part: Final = getattr(openai_responses_api_chunk, "part", None) + _part: Final[PART_UNION_TYPES | Mapping[str, object] | None] = getattr( + openai_responses_api_chunk, "part", None + ) if _part is not None: if isinstance(_part, dict): ResponsesAPIRequestUtils._encode_container_ids_in_annotations( @@ -302,7 +307,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - item: Final = getattr(openai_responses_api_chunk, "item", None) + item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None) if item: encrypted_content: Final = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): @@ -324,9 +329,11 @@ class BaseResponsesAPIStreamingIterator: self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Final[Any | None] = getattr(openai_responses_api_chunk, "response", None) + response_obj: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_obj: - usage_obj: Final[Any | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is not None: try: cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) @@ -414,7 +421,9 @@ class BaseResponsesAPIStreamingIterator: async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ - response_obj: Final = getattr(self.completed_response, "response", None) if self.completed_response else None + response_obj: Final[ResponsesAPIResponse | None] = ( + getattr(self.completed_response, "response", None) if self.completed_response else None + ) error_info: Final = getattr(response_obj, "error", None) if response_obj else None error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) @@ -429,7 +438,7 @@ class BaseResponsesAPIStreamingIterator: def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is None: return try: @@ -506,7 +515,7 @@ class BaseResponsesAPIStreamingIterator: return request_kwargs = getattr(caching_handler, "request_kwargs", None) - if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True: + if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True: return request_kwargs = request_kwargs.copy() preset_cache_key = getattr(caching_handler, "preset_cache_key", None) @@ -606,7 +615,7 @@ class BaseResponsesAPIStreamingIterator: if self.completed_response is None: return - request_payload: Final[dict[str, Any]] = {} + request_payload: Final[dict[str, object]] = {} if isinstance(self.request_data, dict): request_payload.update(self.request_data) try: @@ -695,11 +704,15 @@ class BaseResponsesAPIStreamingIterator: pass -async def call_post_streaming_hooks_for_testing(iterator, chunk): +async def call_post_streaming_hooks_for_testing( + iterator: object, chunk: ResponsesAPIStreamingResponse +) -> ResponsesAPIStreamingResponse: """ Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. """ - hook_fn: Final = getattr(iterator, "_call_post_streaming_deployment_hook", None) + hook_fn: Final[Callable[[ResponsesAPIStreamingResponse], Awaitable[ResponsesAPIStreamingResponse]] | None] = ( + getattr(iterator, "_call_post_streaming_deployment_hook", None) + ) if hook_fn is None: return chunk return await hook_fn(chunk) @@ -1019,7 +1032,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _dump_response_object(obj: Any) -> dict[str, Any]: if hasattr(obj, "model_dump"): return obj.model_dump() - if isinstance(obj, dict): + if _is_json_object(obj): return obj return {} @@ -1684,7 +1697,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + evt_obj: Final[Mapping[str, object]] = json.loads(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -1925,7 +1938,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, Any], + completed_event: dict[str, object], ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into @@ -2065,7 +2078,7 @@ class ManagedResponsesWebSocketHandler: Flat: {"type": "response.create", "input": [...], "model": "...", ...} """ nested: Final = msg_obj.get("response") - response_params: Final[dict[str, Any]] = ( + response_params: Final[dict[str, object]] = ( nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"} ) return { @@ -2076,7 +2089,7 @@ class ManagedResponsesWebSocketHandler: def _apply_history( self, - call_kwargs: dict[str, Any], + call_kwargs: dict[str, object], previous_response_id: str | None, current_messages: list[dict[str, object]], prior_history: list[dict[str, object]], @@ -2129,7 +2142,7 @@ class ManagedResponsesWebSocketHandler: return False return event_provider == self._connection_provider - def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None: + def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a8af4eabb3f..68a6451e273 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3121 + "limit": 3114 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 834 }, "ANN201": { - "limit": 2032 + "limit": 2031 }, "ANN202": { "limit": 865 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1630 + "limit": 1555 }, "ASYNC230": { "limit": 11 @@ -81,7 +81,7 @@ "limit": 1 }, "C901": { - "limit": 315 + "limit": 314 }, "D419": { "limit": 6 @@ -306,7 +306,7 @@ "limit": 0 }, "TID251": { - "limit": 1240 + "limit": 1238 }, "TRY002": { "limit": 528 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0a0cfe9a617..3a670bc7345 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23235 + "limit": 23149 }, "LIT002": { - "limit": 27176 + "limit": 27166 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1091 + "limit": 1086 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16769 + "limit": 16760 }, "LIT011": { "limit": 5598 From f6df762b2537b797b2562af44ac9107ed4fe5c77 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:45:43 -0700 Subject: [PATCH 13/13] test: roll back live router replay membership between tests (#36278) Since #35491, every Router joins the module-global _live_routers weak set at construction, and every model cost map swap replays the deployments of every member on top of the freshly adopted map. #36039 isolated the register_model ledger half of that replay but not this half: under pytest-xdist, a Router created by an earlier test in the same worker that was still referenced (or simply not yet garbage collected) re-registered its deployments during TestPriceDataReloadIntegration::test_distributed_reload_check_function, and register_model hydrated the sparse mocked gpt-3.5-turbo entry into a full ModelInfo dict, failing the exact-equality assert (reruns cannot help since the polluting router survives in the worker process) The autouse isolate_litellm_state fixture now snapshots _live_routers before each test and restores its membership on teardown, so a test's routers stop contributing to cost map rebuilds once the test ends. A canary pair in test_conftest_isolation.py asserts the rollback --- tests/test_litellm/conftest.py | 8 +++++++ tests/test_litellm/test_conftest_isolation.py | 23 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index c0993051d33..0dc8f56f3ce 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -19,6 +19,7 @@ sys.path.insert( import asyncio import litellm +from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.prompt_templates import ( @@ -244,6 +245,8 @@ def isolate_litellm_state(): for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() } + original_live_routers = set(litellm_router_module._live_routers) + # Store LiteLLM logger state. Some tests reconfigure handlers/propagation for # JSON logging and do not restore them, which breaks later caplog-based tests. logger_state = {} @@ -313,6 +316,11 @@ def isolate_litellm_state(): litellm_utils_module._runtime_registered_model_cost.clear() litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + for _router in tuple(litellm_router_module._live_routers): + litellm_router_module._live_routers.discard(_router) + for _router in original_live_routers: + litellm_router_module._live_routers.add(_router) + # Restore logger configuration mutated by logging-focused tests. for logger in ALL_LOGGERS: original_logger_state = logger_state.get(logger.name) diff --git a/tests/test_litellm/test_conftest_isolation.py b/tests/test_litellm/test_conftest_isolation.py index 88889ad7740..15183e68f66 100644 --- a/tests/test_litellm/test_conftest_isolation.py +++ b/tests/test_litellm/test_conftest_isolation.py @@ -1,9 +1,15 @@ import litellm +from litellm import Router +from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module CANARY_MODEL = "conftest-isolation-canary-model" +class _CanaryRouterHolder: + router: Router | None = None + + def test_register_model_ledger_entry_is_scoped_to_this_test(): litellm.register_model({CANARY_MODEL: {"litellm_provider": "openai", "input_cost_per_token": 0.001}}) assert CANARY_MODEL in litellm_utils_module._runtime_registered_model_cost @@ -11,3 +17,20 @@ def test_register_model_ledger_entry_is_scoped_to_this_test(): def test_register_model_ledger_entry_was_rolled_back(): assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost + + +def test_live_router_membership_is_scoped_to_this_test(): + _CanaryRouterHolder.router = Router( + model_list=[ + { + "model_name": "conftest-isolation-canary-router", + "litellm_params": {"model": "openai/conftest-isolation-canary-backend", "api_key": "sk-canary"}, + } + ] + ) + assert _CanaryRouterHolder.router in litellm_router_module._live_routers + + +def test_live_router_membership_was_rolled_back(): + assert _CanaryRouterHolder.router is not None + assert _CanaryRouterHolder.router not in litellm_router_module._live_routers