From 0e2fcde145e9ffb27812badbc314b1c6f2af74d9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 9 Dec 2025 18:17:07 -0800 Subject: [PATCH 001/121] Change health check to read env varlike chat completions --- .../health_endpoints/_health_endpoints.py | 111 +++++++++++++++- .../health_endpoints/test_health_endpoints.py | 118 ++++++++++++++++++ 2 files changed, 223 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 79e9838d115..030843376bf 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -30,9 +30,48 @@ from litellm.proxy.health_check import ( perform_health_check, run_with_timeout, ) +from litellm.secret_managers.main import get_secret #### Health ENDPOINTS #### + +def _resolve_os_environ_variables(params: dict) -> dict: + """ + Resolve os.environ/ environment variables in litellm_params. + + This function recursively processes dictionary values that start with "os.environ/" + by replacing them with the actual environment variable values. + + Args: + params: Dictionary containing litellm_params that may have os.environ/ values + + Returns: + Dictionary with os.environ/ values resolved to actual environment variable values + """ + if not isinstance(params, dict): + return params + + resolved_params = {} + for key, value in params.items(): + if isinstance(value, str) and value.startswith("os.environ/"): + # Resolve the environment variable + resolved_value = get_secret(value) + resolved_params[key] = resolved_value + elif isinstance(value, dict): + # Recursively resolve nested dictionaries + resolved_params[key] = _resolve_os_environ_variables(value) + elif isinstance(value, list): + # Handle lists that might contain dictionaries with os.environ/ values + resolved_params[key] = [ + _resolve_os_environ_variables(item) if isinstance(item, dict) else item + for item in value + ] + else: + resolved_params[key] = value + + return resolved_params + + router = APIRouter() services = Union[ Literal[ @@ -1166,21 +1205,41 @@ async def test_model_connection( Example: ```bash + # If model is configured in proxy_config.yaml, you only need to specify the model name: curl -X POST 'http://localhost:4000/health/test_connection' \\ -H 'Authorization: Bearer sk-1234' \\ -H 'Content-Type: application/json' \\ -d '{ "litellm_params": { - "model": "gpt-4", - "custom_llm_provider": "azure_ai", - "litellm_credential_name": null, - "api_key": "6xxxxxxx", - "api_base": "https://litellm8397336933.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21", + "model": "gpt-4o" + }, + "mode": "chat" + }' + + # The endpoint will automatically use api_key, api_base, etc. from proxy_config.yaml + + # You can also override specific params or test with custom credentials: + curl -X POST 'http://localhost:4000/health/test_connection' \\ + -H 'Authorization: Bearer sk-1234' \\ + -H 'Content-Type: application/json' \\ + -d '{ + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "os.environ/AZURE_OPENAI_API_KEY", + "api_base": "os.environ/AZURE_OPENAI_ENDPOINT", + "api_version": "2024-10-21" }, "mode": "chat" }' ``` + Note: + - If the model is configured in proxy_config.yaml, credentials (api_key, api_base, etc.) + will be automatically loaded from the config (with resolved environment variables). + - You can override specific params by including them in the request. + - You can use `os.environ/VARIABLE_NAME` syntax to reference environment variables, + which will be resolved automatically (same as in proxy_config.yaml). + Returns: dict: A dictionary containing the health check result with either success information or error details. """ @@ -1188,7 +1247,7 @@ async def test_model_connection( from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, ) - from litellm.proxy.proxy_server import premium_user, prisma_client + from litellm.proxy.proxy_server import llm_router, premium_user, prisma_client from litellm.types.router import Deployment, LiteLLM_Params try: @@ -1197,6 +1256,46 @@ async def test_model_connection( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) + + # Get model name from litellm_params + request_litellm_params = litellm_params or {} + model_name = request_litellm_params.get("model") + + # Look up model configuration from router if model name is provided + # This gets the litellm_params from proxy config (with resolved env vars) + config_litellm_params = {} + if model_name and llm_router is not None: + try: + # First try to find by proxy model_name (e.g., "gpt-4o") + deployments = llm_router.get_model_list(model_name=model_name) + + # If not found, try to find by litellm model name (e.g., "azure/gpt-4o") + if not deployments or len(deployments) == 0: + all_deployments = llm_router.get_model_list(model_name=None) + if all_deployments: + for deployment in all_deployments: + if deployment.get("litellm_params", {}).get("model") == model_name: + deployments = [deployment] + break + + if deployments and len(deployments) > 0: + # Use the first deployment's litellm_params as base config + # These already have resolved environment variables from proxy config + config_litellm_params = deployments[0].get("litellm_params", {}).copy() + except Exception as e: + verbose_proxy_logger.debug( + f"Could not find model {model_name} in router: {e}. " + "Proceeding with request params only." + ) + + # Merge: config params (from proxy config) as base, request params override + # This allows users to override specific params while using config for credentials + merged_litellm_params = {**config_litellm_params, **request_litellm_params} + + # Resolve os.environ/ environment variables in any remaining request params + # This handles cases where user explicitly passes os.environ/ values to override config + litellm_params = _resolve_os_environ_variables(merged_litellm_params) + ## Auth check await ModelManagementAuthChecks.can_user_make_model_call( model_params=Deployment( diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 6939a19b7ef..594a92b0dc7 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, db_health_cache, health_services_endpoint, + test_model_connection as health_test_model_connection, ) @@ -126,3 +127,120 @@ async def test_health_services_endpoint_sqs(status, error_message): assert result["message"] == error_message mock_instance.async_health_check.assert_awaited_once() + +@pytest.mark.asyncio +async def test_test_model_connection_loads_config_from_router(): + """ + Test that /health/test_connection automatically loads model configuration + (including resolved environment variables) from the router when model name is provided. + """ + # Mock request + mock_request = MagicMock() + + # Mock user_api_key_dict + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.token = "test-token" + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock router with model configuration + mock_router = MagicMock() + mock_deployment = { + "model_name": "gpt-4o", + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "resolved-api-key-from-env", + "api_base": "https://resolved-endpoint.openai.azure.com/", + "api_version": "2024-10-21", + }, + "model_info": {}, + } + mock_router.get_model_list.return_value = [mock_deployment] + + # Mock ModelManagementAuthChecks - patch at the source module since it's imported inside the function + mock_can_user_make_model_call = AsyncMock() + + # Mock litellm.ahealth_check + mock_health_check_result = { + "status": "healthy", + "response_time_ms": 100, + } + mock_ahealth_check = AsyncMock(return_value=mock_health_check_result) + + # Mock run_with_timeout + mock_run_with_timeout = AsyncMock(return_value=mock_health_check_result) + + # Mock _update_litellm_params_for_health_check + def mock_update_params(model_info, litellm_params): + # Just return params with messages added + params = litellm_params.copy() + params["messages"] = [{"role": "user", "content": "test"}] + return params + + # Mock _resolve_os_environ_variables + def mock_resolve_os_environ(params): + return params + + with patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ), patch( + "litellm.proxy.proxy_server.llm_router", + mock_router, + ), patch( + "litellm.proxy.proxy_server.premium_user", + False, + ), patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + mock_can_user_make_model_call, + ), patch( + "litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check", + mock_ahealth_check, + ), patch( + "litellm.proxy.health_endpoints._health_endpoints.run_with_timeout", + mock_run_with_timeout, + ), patch( + "litellm.proxy.health_endpoints._health_endpoints._update_litellm_params_for_health_check", + mock_update_params, + ), patch( + "litellm.proxy.health_endpoints._health_endpoints._resolve_os_environ_variables", + mock_resolve_os_environ, + ): + # Call the endpoint with only model name (no credentials) + result = await health_test_model_connection( + request=mock_request, + mode="chat", + litellm_params={"model": "gpt-4o"}, + model_info={}, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify router.get_model_list was called with the model name + mock_router.get_model_list.assert_called_once_with(model_name="gpt-4o") + + # Verify that run_with_timeout was called (which wraps ahealth_check) + assert mock_run_with_timeout.called + + # Get the call args to verify merged params + call_args = mock_run_with_timeout.call_args + assert call_args is not None + + # The first arg should be the coroutine from ahealth_check + # We need to check what was passed to ahealth_check + ahealth_check_call_args = mock_ahealth_check.call_args + assert ahealth_check_call_args is not None + model_params = ahealth_check_call_args.kwargs.get("model_params", {}) + + # Verify that config params were loaded and merged + # Note: request params override config params, so model from request is used + assert model_params.get("api_key") == "resolved-api-key-from-env" + assert model_params.get("api_base") == "https://resolved-endpoint.openai.azure.com/" + assert model_params.get("api_version") == "2024-10-21" + assert model_params.get("model") == "gpt-4o" # Request param overrides config param + + # Verify result + assert result["status"] == "success" + assert "result" in result + From 4bcbd8b0a96609af2937083c77584784a309f24b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 12 Dec 2025 18:03:08 -0800 Subject: [PATCH 002/121] Fix callback env variables --- litellm/proxy/proxy_server.py | 11 ++++++++- tests/test_litellm/proxy/test_proxy_server.py | 24 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index da09346503d..bcb47e1d12a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3361,8 +3361,17 @@ class ProxyConfig: decrypted_env_vars = self._decrypt_and_set_db_env_variables( db_param_value, return_original_value=True ) + # Normalize keys when loading from DB so services expecting uppercase + # (e.g. Datadog) can read them even if stored in lowercase. + merged_env_vars: dict = {} + for key, value in decrypted_env_vars.items(): + merged_env_vars[key] = value + upper_key = key.upper() + merged_env_vars[upper_key] = value + os.environ[upper_key] = value + current_config.setdefault("environment_variables", {}).update( - decrypted_env_vars + merged_env_vars ) return current_config elif param_name == "litellm_settings" and isinstance(db_param_value, dict): diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 22a9d5e647b..00c1419a09d 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2609,6 +2609,30 @@ async def test_init_sso_settings_in_db_empty_settings(): assert uppercased_settings == {} +def test_update_config_fields_uppercases_env_vars(monkeypatch): + """ + Ensure environment variables pulled from DB are uppercased when applied so + integrations like Datadog that expect uppercase env keys can read them. + """ + from litellm.proxy.proxy_server import ProxyConfig + + for key in ["DD_API_KEY", "DD_SITE", "dd_api_key", "dd_site"]: + monkeypatch.delenv(key, raising=False) + + proxy_config = ProxyConfig() + updated_config = proxy_config._update_config_fields( + current_config={}, + param_name="environment_variables", + db_param_value={"dd_api_key": "test-api-key", "dd_site": "us5.datadoghq.com"}, + ) + + env_vars = updated_config.get("environment_variables", {}) + assert env_vars["DD_API_KEY"] == "test-api-key" + assert env_vars["DD_SITE"] == "us5.datadoghq.com" + assert os.environ.get("DD_API_KEY") == "test-api-key" + assert os.environ.get("DD_SITE") == "us5.datadoghq.com" + + def test_get_prompt_spec_for_db_prompt_with_versions(): """ Test that _get_prompt_spec_for_db_prompt correctly converts database prompts From 9d420265ef86ef2d708a2df181a3972d262962ad Mon Sep 17 00:00:00 2001 From: Jack Temple Date: Mon, 15 Dec 2025 10:09:05 -0600 Subject: [PATCH 003/121] fix: update UI path handling for non-root Docker and restructure HTML files --- litellm/proxy/proxy_server.py | 87 ++++++++++++++----- tests/test_litellm/proxy/test_proxy_server.py | 27 ++++++ 2 files changed, 94 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index da09346503d..acfe2b42098 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5,6 +5,7 @@ import io import os import random import secrets +import shutil import subprocess import sys import time @@ -939,31 +940,68 @@ origins = ["*"] # get current directory try: current_dir = os.path.dirname(os.path.abspath(__file__)) - ui_path = os.path.join(current_dir, "_experimental", "out") + packaged_ui_path = os.path.join(current_dir, "_experimental", "out") + ui_path = packaged_ui_path litellm_asset_prefix = "/litellm-asset-prefix" - # For non-root Docker, use the pre-built UI from /tmp/litellm_ui - # Support both "true" and "True" for case-insensitive comparison - if os.getenv("LITELLM_NON_ROOT", "").lower() == "true": - non_root_ui_path = "/tmp/litellm_ui" + def _dir_has_content(path: str) -> bool: + try: + return os.path.isdir(path) and any(os.scandir(path)) + except FileNotFoundError: + return False - # Check if the UI was built and exists at the expected location - if os.path.exists(non_root_ui_path) and os.listdir(non_root_ui_path): + # Use a writable runtime UI directory whenever possible. + # This prevents mutating the packaged UI directory (e.g. site-packages or the repo checkout) + # and ensures extensionless routes like /ui/login work via /index.html. + is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true" + runtime_ui_path = "/tmp/litellm_ui" + + if _dir_has_content(runtime_ui_path): + if is_non_root: verbose_proxy_logger.info( - f"Using pre-built UI for non-root Docker: {non_root_ui_path}" + f"Using pre-built UI for non-root Docker: {runtime_ui_path}" ) - verbose_proxy_logger.info( - f"UI files found: {len(os.listdir(non_root_ui_path))} items" - ) - ui_path = non_root_ui_path else: + verbose_proxy_logger.info( + f"Using cached runtime UI directory: {runtime_ui_path}" + ) + ui_path = runtime_ui_path + else: + if is_non_root: verbose_proxy_logger.error( - f"UI not found at {non_root_ui_path}. UI will not be available." + f"UI not found at {runtime_ui_path}. Attempting to populate it from packaged UI." ) verbose_proxy_logger.error( - f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}" + f"Path exists: {os.path.exists(runtime_ui_path)}, Has content: {_dir_has_content(runtime_ui_path)}" ) + try: + os.makedirs(runtime_ui_path, exist_ok=True) + if not _dir_has_content(runtime_ui_path) and _dir_has_content( + packaged_ui_path + ): + shutil.copytree( + packaged_ui_path, + runtime_ui_path, + dirs_exist_ok=True, + ) + except Exception as e: + if is_non_root: + verbose_proxy_logger.exception( + f"Failed to populate runtime UI directory {runtime_ui_path} from {packaged_ui_path}: {e}" + ) + else: + if _dir_has_content(runtime_ui_path): + if is_non_root: + verbose_proxy_logger.info( + f"Using populated UI for non-root Docker: {runtime_ui_path}" + ) + else: + verbose_proxy_logger.info( + f"Using populated runtime UI directory: {runtime_ui_path}" + ) + ui_path = runtime_ui_path + # Only modify files if a custom server root path is set if server_root_path and server_root_path != "/": # Iterate through files in the UI directory @@ -1042,16 +1080,25 @@ try: target_path = os.path.join(target_dir, "index.html") os.makedirs(target_dir, exist_ok=True) - os.replace(file_path, target_path) + try: + os.replace(file_path, target_path) + except FileNotFoundError: + # Another process may have already moved this file. + continue # Handle HTML file restructuring - # Skip this for non-root Docker since it's done at build time - # Support both "true" and "True" for case-insensitive comparison - if os.getenv("LITELLM_NON_ROOT", "").lower() != "true": - _restructure_ui_html_files(ui_path) + # Always restructure the directory we actually serve, but avoid mutating the packaged UI. + # This is critical for extensionless routes like /ui/login (expects login/index.html). + if ui_path != packaged_ui_path: + try: + _restructure_ui_html_files(ui_path) + except PermissionError as e: + verbose_proxy_logger.exception( + f"Permission error while restructuring UI directory {ui_path}: {e}" + ) else: verbose_proxy_logger.info( - "Skipping runtime HTML restructuring for non-root Docker (already done at build time)" + f"Skipping runtime HTML restructuring for packaged UI directory: {ui_path}" ) except Exception: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 22a9d5e647b..c9058616f13 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -15,6 +15,7 @@ import httpx import pytest import yaml from fastapi import FastAPI +from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient sys.path.insert( @@ -196,6 +197,32 @@ def test_restructure_ui_html_files_handles_nested_routes(tmp_path): ) +def test_ui_extensionless_route_requires_restructure(tmp_path): + """Regression for non-root fallback: /ui/login expects login/index.html.""" + + from litellm.proxy import proxy_server + + ui_root = tmp_path / "ui" + ui_root.mkdir() + (ui_root / "index.html").write_text("index") + (ui_root / "login.html").write_text("login") + + fastapi_app = FastAPI() + fastapi_app.mount( + "/ui", StaticFiles(directory=str(ui_root), html=True), name="ui" + ) + client = TestClient(fastapi_app) + + assert client.get("/ui/login.html").status_code == 200 + assert client.get("/ui/login").status_code == 404 + + proxy_server._restructure_ui_html_files(str(ui_root)) + + response = client.get("/ui/login") + assert response.status_code == 200 + assert "login" in response.text + + @pytest.mark.asyncio async def test_initialize_scheduled_jobs_credentials(monkeypatch): """ From 93b1da79118dfc2f1ef30ba523c998f365a3bd69 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Mon, 15 Dec 2025 10:38:54 -0800 Subject: [PATCH 004/121] [Refactor] lazy imports: Use per-attribute lazy imports and extract shared constants (#17994) --- litellm/__init__.py | 35 +++------- litellm/_lazy_imports.py | 86 ++++++++++++++--------- tests/test_litellm/test_lazy_imports.py | 91 +++++++++++++++++++++++++ 3 files changed, 154 insertions(+), 58 deletions(-) create mode 100644 tests/test_litellm/test_lazy_imports.py diff --git a/litellm/__init__.py b/litellm/__init__.py index ef44aa53a13..05968e8a7c0 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1563,41 +1563,24 @@ if TYPE_CHECKING: def __getattr__(name: str) -> Any: """Lazy import handler for cost_calculator and litellm_logging functions.""" - # Lazy load cost_calculator functions - _cost_calculator_names = ( - "completion_cost", - "cost_per_token", - "response_cost_calculator", + from ._lazy_imports import ( + COST_CALCULATOR_NAMES, + LITELLM_LOGGING_NAMES, + UTILS_NAMES, ) - if name in _cost_calculator_names: + + # Lazy load cost_calculator functions + if name in COST_CALCULATOR_NAMES: from ._lazy_imports import _lazy_import_cost_calculator return _lazy_import_cost_calculator(name) # Lazy load litellm_logging functions - _litellm_logging_names = ( - "Logging", - "modify_integration", - ) - if name in _litellm_logging_names: + if name in LITELLM_LOGGING_NAMES: from ._lazy_imports import _lazy_import_litellm_logging return _lazy_import_litellm_logging(name) # Lazy load utils functions - _utils_names = ( - "exception_type", "get_optional_params", "get_response_string", "token_counter", - "create_pretrained_tokenizer", "create_tokenizer", "supports_function_calling", - "supports_web_search", "supports_url_context", "supports_response_schema", - "supports_parallel_function_calling", "supports_vision", "supports_audio_input", - "supports_audio_output", "supports_system_messages", "supports_reasoning", - "get_litellm_params", "acreate", "get_max_tokens", "get_model_info", - "register_prompt_template", "validate_environment", "check_valid_key", - "register_model", "encode", "decode", "_calculate_retry_after", "_should_retry", - "get_supported_openai_params", "get_api_base", "get_first_chars_messages", - "ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse", - "TranscriptionResponse", "TextCompletionResponse", "get_provider_fields", - "ModelResponseListIterator", "get_valid_models", - ) - if name in _utils_names: + if name in UTILS_NAMES: from ._lazy_imports import _lazy_import_utils return _lazy_import_utils(name) diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 91b16864de1..b87b9c955fb 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -5,6 +5,35 @@ def _get_litellm_globals() -> dict: """Helper to get the globals dictionary of the litellm module.""" return sys.modules["litellm"].__dict__ +# Cost calculator names that support lazy loading via _lazy_import_cost_calculator +COST_CALCULATOR_NAMES = ( + "completion_cost", + "cost_per_token", + "response_cost_calculator", +) + +# Litellm logging names that support lazy loading via _lazy_import_litellm_logging +LITELLM_LOGGING_NAMES = ( + "Logging", + "modify_integration", +) + +# Utils names that support lazy loading via _lazy_import_utils +UTILS_NAMES = ( + "exception_type", "get_optional_params", "get_response_string", "token_counter", + "create_pretrained_tokenizer", "create_tokenizer", "supports_function_calling", + "supports_web_search", "supports_url_context", "supports_response_schema", + "supports_parallel_function_calling", "supports_vision", "supports_audio_input", + "supports_audio_output", "supports_system_messages", "supports_reasoning", + "get_litellm_params", "acreate", "get_max_tokens", "get_model_info", + "register_prompt_template", "validate_environment", "check_valid_key", + "register_model", "encode", "decode", "_calculate_retry_after", "_should_retry", + "get_supported_openai_params", "get_api_base", "get_first_chars_messages", + "ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse", + "TranscriptionResponse", "TextCompletionResponse", "get_provider_fields", + "ModelResponseListIterator", "get_valid_models", +) + # Lazy import for utils module - imports only the requested item by name. # Note: PLR0915 (too many statements) is suppressed because the many if statements # are intentional - each attribute is imported individually only when requested, @@ -218,42 +247,35 @@ def _lazy_import_utils(name: str) -> Any: # noqa: PLR0915 def _lazy_import_cost_calculator(name: str) -> Any: """Lazy import for cost_calculator functions.""" _globals = _get_litellm_globals() - from .cost_calculator import ( - completion_cost as _completion_cost, - cost_per_token as _cost_per_token, - response_cost_calculator as _response_cost_calculator, - ) + if name == "completion_cost": + from .cost_calculator import completion_cost as _completion_cost + _globals["completion_cost"] = _completion_cost + return _completion_cost - _cost_functions = { - "completion_cost": _completion_cost, - "cost_per_token": _cost_per_token, - "response_cost_calculator": _response_cost_calculator, - } + if name == "cost_per_token": + from .cost_calculator import cost_per_token as _cost_per_token + _globals["cost_per_token"] = _cost_per_token + return _cost_per_token - func = _cost_functions[name] - _globals[name] = func - return func + if name == "response_cost_calculator": + from .cost_calculator import response_cost_calculator as _response_cost_calculator + _globals["response_cost_calculator"] = _response_cost_calculator + return _response_cost_calculator + + raise AttributeError(f"Cost calculator lazy import: unknown attribute {name!r}") def _lazy_import_litellm_logging(name: str) -> Any: """Lazy import for litellm_logging module.""" _globals = _get_litellm_globals() - try: - from litellm.litellm_core_utils.litellm_logging import ( - Logging as _Logging, - modify_integration as _modify_integration, - ) - - _logging_objects = { - "Logging": _Logging, - "modify_integration": _modify_integration, - } - - obj = _logging_objects[name] - _globals[name] = obj - return obj - except Exception as e: - raise AttributeError( - f"module 'litellm' has no attribute {name!r}. " - f"Lazy import failed: {e}" - ) from e \ No newline at end of file + if name == "Logging": + from litellm.litellm_core_utils.litellm_logging import Logging as _Logging + _globals["Logging"] = _Logging + return _Logging + + if name == "modify_integration": + from litellm.litellm_core_utils.litellm_logging import modify_integration as _modify_integration + _globals["modify_integration"] = _modify_integration + return _modify_integration + + raise AttributeError(f"Litellm logging lazy import: unknown attribute {name!r}") \ No newline at end of file diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py new file mode 100644 index 00000000000..e36acf22eb8 --- /dev/null +++ b/tests/test_litellm/test_lazy_imports.py @@ -0,0 +1,91 @@ +"""Simple tests for lazy import functionality.""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm._lazy_imports import ( + COST_CALCULATOR_NAMES, + LITELLM_LOGGING_NAMES, + UTILS_NAMES, + _lazy_import_cost_calculator, + _lazy_import_litellm_logging, + _lazy_import_utils, +) + + +def _clear_names_from_globals(names: tuple): + """Clear all names from litellm globals.""" + for name in names: + if name in litellm.__dict__: + del litellm.__dict__[name] + + +def _verify_only_requested_name_imported(name: str, all_names: tuple): + """Verify that only the requested name is in globals, not the others.""" + for other_name in all_names: + if other_name != name: + assert other_name not in litellm.__dict__, f"{other_name} should not be imported when importing {name}" + + +def test_cost_calculator_lazy_imports(): + """Test that all cost calculator functions can be lazy imported.""" + # Test each name individually - only that name should be imported + for name in COST_CALCULATOR_NAMES: + # Clear all names before importing just one + _clear_names_from_globals(COST_CALCULATOR_NAMES) + + func = _lazy_import_cost_calculator(name) + assert func is not None + assert callable(func) + assert name in litellm.__dict__ + + # Verify only the requested name is in globals, not the others + _verify_only_requested_name_imported(name, COST_CALCULATOR_NAMES) + + +def test_litellm_logging_lazy_imports(): + """Test that all litellm_logging items can be lazy imported.""" + # Test each name individually - only that name should be imported + for name in LITELLM_LOGGING_NAMES: + # Clear all names before importing just one + _clear_names_from_globals(LITELLM_LOGGING_NAMES) + + item = _lazy_import_litellm_logging(name) + assert item is not None + assert name in litellm.__dict__ + + # Verify only the requested name is in globals, not the others + _verify_only_requested_name_imported(name, LITELLM_LOGGING_NAMES) + + +def test_utils_lazy_imports(): + """Test that all utils functions can be lazy imported.""" + # Test each name individually - only that name should be imported + for name in UTILS_NAMES: + # Clear all names before importing just one + _clear_names_from_globals(UTILS_NAMES) + + attr = _lazy_import_utils(name) + assert attr is not None + assert name in litellm.__dict__ + + # Verify only the requested name is in globals, not the others + _verify_only_requested_name_imported(name, UTILS_NAMES) + + +def test_unknown_attribute_raises_error(): + """Test that unknown attributes raise AttributeError.""" + with pytest.raises(AttributeError): + _lazy_import_cost_calculator("unknown") + + with pytest.raises(AttributeError): + _lazy_import_litellm_logging("unknown") + + with pytest.raises(AttributeError): + _lazy_import_utils("unknown") + From 0629dcfdd5f10e5ab8733c97f3909019ba7a3a6e Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Mon, 15 Dec 2025 11:50:22 -0800 Subject: [PATCH 005/121] [Refactor] `litellm/init.py`: lazy load http handlers (#17997) --- litellm/__init__.py | 19 +++++++---- litellm/_lazy_imports.py | 42 +++++++++++++++++++++++-- tests/test_litellm/test_lazy_imports.py | 14 +++++++++ 3 files changed, 67 insertions(+), 8 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 05968e8a7c0..bd7b03064dd 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -26,7 +26,6 @@ from typing import ( ) from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams from litellm.types.integrations.datadog import DatadogInitParams -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.caching.caching import Cache, DualCache, RedisCache, InMemoryCache from litellm.caching.llm_caching_handler import LLMClientCache from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES @@ -422,10 +421,6 @@ disable_aiohttp_trust_env: bool = ( force_ipv4: bool = ( False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. ) -module_level_aclient = AsyncHTTPHandler( - timeout=request_timeout, client_alias="module level aclient" -) -module_level_client = HTTPHandler(timeout=request_timeout) #### RETRIES #### num_retries: Optional[int] = None # per model endpoint @@ -1520,6 +1515,7 @@ def set_global_gitlab_config(config: Dict[str, Any]) -> None: if TYPE_CHECKING: from litellm.types.utils import ModelInfo as _ModelInfoType + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler # Cost calculator functions cost_per_token: Callable[..., Tuple[float, float]] @@ -1560,13 +1556,18 @@ if TYPE_CHECKING: # Response types - truly lazy loaded only (not in main.py or elsewhere) ModelResponseListIterator: Type[Any] + # HTTP handler singletons (created lazily via __getattr__ at runtime) + module_level_aclient: AsyncHTTPHandler + module_level_client: HTTPHandler + def __getattr__(name: str) -> Any: - """Lazy import handler for cost_calculator and litellm_logging functions.""" + """Lazy import handler""" from ._lazy_imports import ( COST_CALCULATOR_NAMES, LITELLM_LOGGING_NAMES, UTILS_NAMES, + HTTP_HANDLER_NAMES, ) # Lazy load cost_calculator functions @@ -1584,6 +1585,12 @@ def __getattr__(name: str) -> Any: from ._lazy_imports import _lazy_import_utils return _lazy_import_utils(name) + # Lazy-load HTTP handler singletons used across the codebase + if name in HTTP_HANDLER_NAMES: + from ._lazy_imports import _lazy_import_http_handlers + + return _lazy_import_http_handlers(name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index b87b9c955fb..17772682599 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, cast import sys def _get_litellm_globals() -> dict: @@ -34,6 +34,12 @@ UTILS_NAMES = ( "ModelResponseListIterator", "get_valid_models", ) +# HTTP handler names that support lazy loading via _lazy_import_http_handlers +HTTP_HANDLER_NAMES = ( + "module_level_aclient", + "module_level_client", +) + # Lazy import for utils module - imports only the requested item by name. # Note: PLR0915 (too many statements) is suppressed because the many if statements # are intentional - each attribute is imported individually only when requested, @@ -278,4 +284,36 @@ def _lazy_import_litellm_logging(name: str) -> Any: _globals["modify_integration"] = _modify_integration return _modify_integration - raise AttributeError(f"Litellm logging lazy import: unknown attribute {name!r}") \ No newline at end of file + raise AttributeError(f"Litellm logging lazy import: unknown attribute {name!r}") + + +def _lazy_import_http_handlers(name: str) -> Any: + """Lazy import and instantiate module-level HTTP handlers.""" + _globals = _get_litellm_globals() + + if name == "module_level_aclient": + # Use shared async client factory instead of directly instantiating AsyncHTTPHandler + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + + timeout = _globals.get("request_timeout") + params = {"timeout": timeout, "client_alias": "module level aclient"} + # llm_provider is only used for cache keying; use a string identifier but + # cast to Any so static type checkers don't complain about the literal. + provider_id = cast(Any, "litellm_module_level_client") + async_client = get_async_httpx_client( + llm_provider=provider_id, + params=params, + ) + _globals["module_level_aclient"] = async_client + return async_client + + if name == "module_level_client": + # Import handler type locally to avoid heavy imports at module load time + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + timeout = _globals.get("request_timeout") + sync_client = HTTPHandler(timeout=timeout) + _globals["module_level_client"] = sync_client + return sync_client + + raise AttributeError(f"HTTP handlers lazy import: unknown attribute {name!r}") \ No newline at end of file diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index e36acf22eb8..42060579842 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -12,9 +12,11 @@ from litellm._lazy_imports import ( COST_CALCULATOR_NAMES, LITELLM_LOGGING_NAMES, UTILS_NAMES, + HTTP_HANDLER_NAMES, _lazy_import_cost_calculator, _lazy_import_litellm_logging, _lazy_import_utils, + _lazy_import_http_handlers, ) @@ -78,6 +80,18 @@ def test_utils_lazy_imports(): _verify_only_requested_name_imported(name, UTILS_NAMES) +def test_http_handler_lazy_imports(): + """Test that HTTP handler singletons can be lazy imported.""" + for name in HTTP_HANDLER_NAMES: + _clear_names_from_globals(HTTP_HANDLER_NAMES) + + handler = _lazy_import_http_handlers(name) + assert handler is not None + assert name in litellm.__dict__ + + _verify_only_requested_name_imported(name, HTTP_HANDLER_NAMES) + + def test_unknown_attribute_raises_error(): """Test that unknown attributes raise AttributeError.""" with pytest.raises(AttributeError): From 8f647dd25bd6f1c3b60a6f71acb006d40275178a Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Mon, 15 Dec 2025 12:13:51 -0800 Subject: [PATCH 006/121] [Refactor] litellm/init.py: lazy load caches (#18001) --- litellm/__init__.py | 10 +++++-- litellm/_lazy_imports.py | 39 +++++++++++++++++++++++++ tests/test_litellm/test_lazy_imports.py | 20 +++++++++++++ 3 files changed, 67 insertions(+), 2 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index bd7b03064dd..67fd8a7b075 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -26,7 +26,6 @@ from typing import ( ) from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams from litellm.types.integrations.datadog import DatadogInitParams -from litellm.caching.caching import Cache, DualCache, RedisCache, InMemoryCache from litellm.caching.llm_caching_handler import LLMClientCache from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES from litellm.types.utils import ( @@ -332,7 +331,7 @@ caching: bool = ( caching_with_models: bool = ( False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 ) -cache: Optional[Cache] = ( +cache: Optional["Cache"] = ( None # cache object <- use this - https://docs.litellm.ai/docs/caching ) default_in_memory_ttl: Optional[float] = None @@ -1516,6 +1515,7 @@ def set_global_gitlab_config(config: Dict[str, Any]) -> None: if TYPE_CHECKING: from litellm.types.utils import ModelInfo as _ModelInfoType from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.caching.caching import Cache # Cost calculator functions cost_per_token: Callable[..., Tuple[float, float]] @@ -1567,6 +1567,7 @@ def __getattr__(name: str) -> Any: COST_CALCULATOR_NAMES, LITELLM_LOGGING_NAMES, UTILS_NAMES, + CACHING_NAMES, HTTP_HANDLER_NAMES, ) @@ -1585,6 +1586,11 @@ def __getattr__(name: str) -> Any: from ._lazy_imports import _lazy_import_utils return _lazy_import_utils(name) + # Lazy load caching classes + if name in CACHING_NAMES: + from ._lazy_imports import _lazy_import_caching + return _lazy_import_caching(name) + # Lazy-load HTTP handler singletons used across the codebase if name in HTTP_HANDLER_NAMES: from ._lazy_imports import _lazy_import_http_handlers diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 17772682599..c9655f0d2ff 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -34,6 +34,14 @@ UTILS_NAMES = ( "ModelResponseListIterator", "get_valid_models", ) +# Caching / cache classes that support lazy loading via _lazy_import_caching +CACHING_NAMES = ( + "Cache", + "DualCache", + "RedisCache", + "InMemoryCache", +) + # HTTP handler names that support lazy loading via _lazy_import_http_handlers HTTP_HANDLER_NAMES = ( "module_level_aclient", @@ -271,6 +279,37 @@ def _lazy_import_cost_calculator(name: str) -> Any: raise AttributeError(f"Cost calculator lazy import: unknown attribute {name!r}") +def _lazy_import_caching(name: str) -> Any: + """Lazy import for caching module classes.""" + _globals = _get_litellm_globals() + + if name == "Cache": + from litellm.caching.caching import Cache as _Cache + + _globals["Cache"] = _Cache + return _Cache + + if name == "DualCache": + from litellm.caching.caching import DualCache as _DualCache + + _globals["DualCache"] = _DualCache + return _DualCache + + if name == "RedisCache": + from litellm.caching.caching import RedisCache as _RedisCache + + _globals["RedisCache"] = _RedisCache + return _RedisCache + + if name == "InMemoryCache": + from litellm.caching.caching import InMemoryCache as _InMemoryCache + + _globals["InMemoryCache"] = _InMemoryCache + return _InMemoryCache + + raise AttributeError(f"Caching lazy import: unknown attribute {name!r}") + + def _lazy_import_litellm_logging(name: str) -> Any: """Lazy import for litellm_logging module.""" _globals = _get_litellm_globals() diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index 42060579842..c623ecc4b0e 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -12,10 +12,12 @@ from litellm._lazy_imports import ( COST_CALCULATOR_NAMES, LITELLM_LOGGING_NAMES, UTILS_NAMES, + CACHING_NAMES, HTTP_HANDLER_NAMES, _lazy_import_cost_calculator, _lazy_import_litellm_logging, _lazy_import_utils, + _lazy_import_caching, _lazy_import_http_handlers, ) @@ -80,6 +82,21 @@ def test_utils_lazy_imports(): _verify_only_requested_name_imported(name, UTILS_NAMES) +def test_caching_lazy_imports(): + """Test that all caching classes can be lazy imported.""" + # Test each name individually - only that name should be imported + for name in CACHING_NAMES: + # Clear all names before importing just one + _clear_names_from_globals(CACHING_NAMES) + + cls = _lazy_import_caching(name) + assert cls is not None + assert name in litellm.__dict__ + + # Verify only the requested name is in globals, not the others + _verify_only_requested_name_imported(name, CACHING_NAMES) + + def test_http_handler_lazy_imports(): """Test that HTTP handler singletons can be lazy imported.""" for name in HTTP_HANDLER_NAMES: @@ -103,3 +120,6 @@ def test_unknown_attribute_raises_error(): with pytest.raises(AttributeError): _lazy_import_utils("unknown") + with pytest.raises(AttributeError): + _lazy_import_caching("unknown") + From 764a31f6245509cf732e695cc47a59e70b794e0f Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Mon, 15 Dec 2025 12:33:01 -0800 Subject: [PATCH 007/121] refactor: lazy load get_modified_max_tokens (#18002) --- litellm/__init__.py | 7 ++++++- litellm/_lazy_imports.py | 20 ++++++++++++++++++++ tests/test_litellm/test_lazy_imports.py | 17 +++++++++++++++++ 3 files changed, 43 insertions(+), 1 deletion(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 67fd8a7b075..80625e0b189 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1065,7 +1065,6 @@ openai_video_generation_models = ["sora-2"] from .timeout import timeout from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.core_helpers import remove_index_from_tool_calls -from litellm.litellm_core_utils.token_counter import get_modified_max_tokens # client must be imported immediately as it's used as a decorator at function definition time from .utils import client # Note: Most other utils imports are lazy-loaded via __getattr__ to avoid loading utils.py @@ -1567,6 +1566,7 @@ def __getattr__(name: str) -> Any: COST_CALCULATOR_NAMES, LITELLM_LOGGING_NAMES, UTILS_NAMES, + TOKEN_COUNTER_NAMES, CACHING_NAMES, HTTP_HANDLER_NAMES, ) @@ -1586,6 +1586,11 @@ def __getattr__(name: str) -> Any: from ._lazy_imports import _lazy_import_utils return _lazy_import_utils(name) + # Lazy load token counter utilities + if name in TOKEN_COUNTER_NAMES: + from ._lazy_imports import _lazy_import_token_counter + return _lazy_import_token_counter(name) + # Lazy load caching classes if name in CACHING_NAMES: from ._lazy_imports import _lazy_import_caching diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index c9655f0d2ff..1fbf3f1be0f 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -34,6 +34,11 @@ UTILS_NAMES = ( "ModelResponseListIterator", "get_valid_models", ) +# Token counter names that support lazy loading via _lazy_import_token_counter +TOKEN_COUNTER_NAMES = ( + "get_modified_max_tokens", +) + # Caching / cache classes that support lazy loading via _lazy_import_caching CACHING_NAMES = ( "Cache", @@ -279,6 +284,21 @@ def _lazy_import_cost_calculator(name: str) -> Any: raise AttributeError(f"Cost calculator lazy import: unknown attribute {name!r}") +def _lazy_import_token_counter(name: str) -> Any: + """Lazy import for token_counter utilities.""" + _globals = _get_litellm_globals() + + if name == "get_modified_max_tokens": + from litellm.litellm_core_utils.token_counter import ( + get_modified_max_tokens as _get_modified_max_tokens, + ) + + _globals["get_modified_max_tokens"] = _get_modified_max_tokens + return _get_modified_max_tokens + + raise AttributeError(f"Token counter lazy import: unknown attribute {name!r}") + + def _lazy_import_caching(name: str) -> Any: """Lazy import for caching module classes.""" _globals = _get_litellm_globals() diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index c623ecc4b0e..08737df4e4a 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -12,11 +12,13 @@ from litellm._lazy_imports import ( COST_CALCULATOR_NAMES, LITELLM_LOGGING_NAMES, UTILS_NAMES, + TOKEN_COUNTER_NAMES, CACHING_NAMES, HTTP_HANDLER_NAMES, _lazy_import_cost_calculator, _lazy_import_litellm_logging, _lazy_import_utils, + _lazy_import_token_counter, _lazy_import_caching, _lazy_import_http_handlers, ) @@ -97,6 +99,18 @@ def test_caching_lazy_imports(): _verify_only_requested_name_imported(name, CACHING_NAMES) +def test_token_counter_lazy_imports(): + """Test that token counter utilities can be lazy imported.""" + for name in TOKEN_COUNTER_NAMES: + _clear_names_from_globals(TOKEN_COUNTER_NAMES) + + func = _lazy_import_token_counter(name) + assert func is not None + assert name in litellm.__dict__ + + _verify_only_requested_name_imported(name, TOKEN_COUNTER_NAMES) + + def test_http_handler_lazy_imports(): """Test that HTTP handler singletons can be lazy imported.""" for name in HTTP_HANDLER_NAMES: @@ -123,3 +137,6 @@ def test_unknown_attribute_raises_error(): with pytest.raises(AttributeError): _lazy_import_caching("unknown") + with pytest.raises(AttributeError): + _lazy_import_token_counter("unknown") + From d7e0044118eadccf8f67e90306cd9492cb79779c Mon Sep 17 00:00:00 2001 From: vasilisazayka Date: Tue, 16 Dec 2025 01:25:26 +0400 Subject: [PATCH 008/121] [docs] update SAP docs (#17974) * docs(sap): update documentation * docs(sap): update documentation * docs(sap): update documentation * docs(sap): update documentation --- docs/my-website/docs/providers/sap.md | 65 ++++++++++++++++++++++----- 1 file changed, 53 insertions(+), 12 deletions(-) diff --git a/docs/my-website/docs/providers/sap.md b/docs/my-website/docs/providers/sap.md index a9183b9c0df..4bc72c27045 100644 --- a/docs/my-website/docs/providers/sap.md +++ b/docs/my-website/docs/providers/sap.md @@ -5,12 +5,12 @@ import TabItem from '@theme/TabItem'; LiteLLM supports SAP Generative AI Hub's Orchestration Service. -| Property | Details | -|-------|-------| -| Description | SAP's Generative AI Hub provides access to foundation models through the AI Core orchestration service. | -| Provider Route on LiteLLM | `sap/` | -| Supported Endpoints | `/chat/completions` | -| API Reference | [SAP AI Core Documentation](https://help.sap.com/docs/sap-ai-core) | +| Property | Details | +|-------|--------------------------------------------------------------------------------------------------------------------------------------------------------| +| Description | SAP's Generative AI Hub provides access to OpenAI, Anthropic, Gemini, Mistral, NVIDIA, Amazon, and SAP LLMs through the AI Core orchestration service. | +| Provider Route on LiteLLM | `sap/` | +| Supported Endpoints | `/chat/completions`, `/embeddings` | +| API Reference | [SAP AI Core Documentation](https://help.sap.com/docs/sap-ai-core) | ## Authentication @@ -23,7 +23,14 @@ SAP Generative AI Hub uses service key authentication. You can provide credentia import os os.environ["AICORE_SERVICE_KEY"] = '{"clientid": "...", "clientsecret": "...", ...}' ``` - +3. **Environment variables** - Set the following list of credentials in .env file +
+AICORE_AUTH_URL = "https://* * * .authentication.sap.hana.ondemand.com/oauth/token",
+AICORE_CLIENT_ID  = " *** ",
+AICORE_CLIENT_SECRET = " *** ",
+AICORE_RESOURCE_GROUP = " *** ",
+AICORE_BASE_URL = "https://api.ai.***.cfapps.sap.hana.ondemand.com/v2"
+
## Usage - LiteLLM Python SDK ```python showLineNumbers title="SAP Chat Completion" @@ -55,16 +62,33 @@ for chunk in response: print(chunk.choices[0].delta.content or "", end="") ``` +```python showLineNumbers title="SAP Embedding" +from litellm import embedding +import os + +os.environ["AICORE_SERVICE_KEY"] = '{"clientid": "...", "clientsecret": "...", ...}' + +result = embedding( + model="sap/text-embedding-3-small", + input="Answer to the ultimate question of life, the universe, and everything is 42") +print(result.data[0]) +``` + ## Usage - LiteLLM Proxy Add to your LiteLLM Proxy config: ```yaml showLineNumbers title="config.yaml" model_list: - - model_name: sap-gpt4 + - model_name: "sap/*" litellm_params: - model: sap/gpt-4 - api_key: os.environ/AICORE_SERVICE_KEY + model: "sap/*" + +general_settings: + master_key: your-proxy-api-key + +environment_variables: + AICORE_SERVICE_KEY: '{"clientid": "...", "clientsecret": "...", ...}' ``` Start the proxy: @@ -81,7 +105,7 @@ curl http://localhost:4000/v1/chat/completions \ -H "Content-Type: application/json" \ -H "Authorization: Bearer your-proxy-api-key" \ -d '{ - "model": "sap-gpt4", + "model": "sap/gpt-4", "messages": [{"role": "user", "content": "Hello"}] }' ``` @@ -98,12 +122,29 @@ client = OpenAI( ) response = client.chat.completions.create( - model="sap-gpt4", + model="sap/gpt-4", messages=[{"role": "user", "content": "Hello"}] ) print(response.choices[0].message.content) ``` + + + +```python showLineNumbers title="LiteLLM SDK" +import os +import litellm +os.environ["LITELLM_PROXY_API_KEY"] = "your-proxy-api-key" +litellm.use_litellm_proxy = True # it is important to set this parameter +response = litellm.completion( + model="sap/gpt-4o", + messages=[{ "content": "Hello, how are you?","role": "user"}], + api_base="http://your-proxy-api-base" +) + +print(response) +``` + From b57b1beb612607df1bcba427788cee0a2cb038d9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 15 Dec 2025 13:26:20 -0800 Subject: [PATCH 009/121] [Feat] Guardrails - litellm content filter (#18007) * add br ssn * better description --- ...odel_prices_and_context_window_backup.json | 2 +- .../litellm_content_filter/patterns.json | 56 +++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8da7b93699e..4b016bc6ca6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22282,7 +22282,7 @@ "supports_reasoning": true, "supports_tool_choice": true }, - "openrouter/openai/gpt-5.2": { + "openrouter/openai/gpt-5.2": { "input_cost_per_image": 0, "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json index 193d0868072..b87bd397aad 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json @@ -317,6 +317,62 @@ "category": "PII Patterns", "action": "MASK", "description": "Detects Dutch BSN numbers with contextual keywords" + }, + { + "name": "br_cpf", + "display_name": "CPF - Brazilian Tax ID / Social Security (Formatted)", + "pattern": "\\b\\d{3}\\.\\d{3}\\.\\d{3}-\\d{2}\\b", + "category": "Brazilian PII Patterns", + "description": "Detects Brazilian CPF numbers (XXX.XXX.XXX-XX format)" + }, + { + "name": "br_cpf_no_format", + "display_name": "CPF - Brazilian Tax ID / Social Security (Unformatted)", + "pattern": "\\b\\d{11}\\b", + "category": "Brazilian PII Patterns", + "description": "Detects Brazilian CPF numbers without formatting (11 digits)" + }, + { + "name": "br_phone", + "display_name": "Brazilian Phone Number (Landline & Mobile)", + "pattern": "\\b(?:\\+?55[\\s.-]?)?\\(?([1-9]{2})\\)?[\\s.-]?(?:[2-9]\\d{3,4})[\\s.-]?(\\d{4})\\b", + "category": "Brazilian PII Patterns", + "description": "Detects Brazilian phone numbers with various area codes (landline and mobile)" + }, + { + "name": "br_phone_mobile", + "display_name": "Brazilian Mobile Phone Number", + "pattern": "\\b(?:\\+?55[\\s.-]?)?\\(?([1-9]{2})\\)?[\\s.-]?9[\\s.-]?\\d{4}[\\s.-]?\\d{4}\\b", + "category": "Brazilian PII Patterns", + "description": "Detects Brazilian mobile phone numbers (9 prefix for mobile)" + }, + { + "name": "br_cep", + "display_name": "CEP - Brazilian Zip / Postal Code", + "pattern": "\\b\\d{5}-\\d{3}\\b", + "category": "Brazilian PII Patterns", + "description": "Detects Brazilian CEP postal codes (XXXXX-XXX format)" + }, + { + "name": "br_address", + "display_name": "Brazilian Street Address", + "pattern": "\\b(?:Rua|Avenida|Av\\.|R\\.|Travessa|Alameda|Praça|Rodovia)\\s+[A-Za-zÀ-ÿ\\s]+,?\\s*(?:n[°º]?|número)?\\s*\\d+", + "category": "Brazilian PII Patterns", + "description": "Detects Brazilian street addresses with common prefixes" + }, + { + "name": "br_cnpj", + "display_name": "CNPJ - Brazilian Company Tax ID", + "pattern": "\\b\\d{2}\\.\\d{3}\\.\\d{3}/\\d{4}-\\d{2}\\b", + "category": "Brazilian PII Patterns", + "description": "Detects Brazilian CNPJ company registration numbers (XX.XXX.XXX/XXXX-XX format)" + }, + { + "name": "br_rg", + "display_name": "RG - Brazilian National Identity Card", + "pattern": "\\b\\d{1,2}\\.?\\d{3}\\.?\\d{3}-?[0-9Xx]\\b", + "category": "Brazilian PII Patterns", + "description": "Detects Brazilian RG identity card numbers" } ] } From defea8b8876c38474a7decac2a89ff5b20fd8b41 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Tue, 16 Dec 2025 07:02:07 +0900 Subject: [PATCH 010/121] fix: mcp deepcopy error --- .../mcp_server/ui_session_utils.py | 4 +-- .../mcp_server/test_ui_session_utils.py | 28 ++++++++++++++++++- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py index 6572b831a27..37a3228ebf0 100644 --- a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py +++ b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py @@ -16,9 +16,9 @@ def clone_user_api_key_auth_with_team( """Return a deep copy of the auth context with a different team id.""" try: - cloned_auth = user_api_key_auth.model_copy(deep=True) + cloned_auth = user_api_key_auth.model_copy() except AttributeError: - cloned_auth = user_api_key_auth.copy(deep=True) # type: ignore[attr-defined] + cloned_auth = user_api_key_auth.copy() # type: ignore[attr-defined] cloned_auth.team_id = team_id return cloned_auth diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py index f372f7b181c..35cfbee0d54 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py @@ -1,7 +1,9 @@ -import pytest +import threading from types import SimpleNamespace from unittest.mock import AsyncMock +import pytest + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import UserAPIKeyAuth @@ -90,3 +92,27 @@ async def test_build_effective_auth_contexts_returns_original_when_no_resolution assert contexts == [user_auth] mock_resolve.assert_awaited_once_with(user_auth) + +@pytest.mark.asyncio +async def test_build_effective_auth_contexts_handles_unpicklable_parent_span(monkeypatch): + class DummySpan: + def __init__(self) -> None: + self._lock = threading.RLock() + + parent_span = DummySpan() + user_auth = UserAPIKeyAuth( + team_id=UI_SESSION_TOKEN_TEAM_ID, + user_id="user-span", + parent_otel_span=parent_span, + ) + + mock_resolve = AsyncMock(return_value=["team-span"]) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.ui_session_utils.resolve_ui_session_team_ids", + mock_resolve, + ) + + contexts = await build_effective_auth_contexts(user_auth) + + assert contexts[0].team_id == "team-span" + assert contexts[0].parent_otel_span is parent_span From df19a747a2bb30970aac959d617a7efb56522ea6 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Mon, 15 Dec 2025 19:16:36 -0300 Subject: [PATCH 011/121] feat(custom_llm): add image_edit and aimage_edit support (#17999) * feat(custom_llm): add image_edit and aimage_edit support Add support for image_edit and aimage_edit methods in CustomLLM class, allowing users to implement custom image editing providers. Changes: - Add image_edit() and aimage_edit() methods to CustomLLM base class - Add custom provider detection in litellm.image_edit() function - Add tests for sync and async image_edit with custom handlers * docs: add image_edit to CustomLLM documentation - Add /v1/images/edits to supported routes - Add Image Edit section with example - Update Custom Handler Spec with image_edit methods --- .../docs/providers/custom_llm_server.md | 108 ++++++++++++++++++ litellm/images/main.py | 53 +++++++++ litellm/llms/custom_llm.py | 30 +++++ tests/local_testing/test_custom_llm.py | 101 ++++++++++++++++ 4 files changed, 292 insertions(+) diff --git a/docs/my-website/docs/providers/custom_llm_server.md b/docs/my-website/docs/providers/custom_llm_server.md index 61099d1a358..4fcbf8942ce 100644 --- a/docs/my-website/docs/providers/custom_llm_server.md +++ b/docs/my-website/docs/providers/custom_llm_server.md @@ -17,6 +17,7 @@ Supported Routes: - `/v1/completions` -> `litellm.atext_completion` - `/v1/embeddings` -> `litellm.aembedding` - `/v1/images/generations` -> `litellm.aimage_generation` +- `/v1/images/edits` -> `litellm.aimage_edit` - `/v1/messages` -> `litellm.acompletion` @@ -263,6 +264,83 @@ Expected Response } ``` +## Image Edit + +1. Setup your `custom_handler.py` file +```python +import litellm +from litellm import CustomLLM +from litellm.types.utils import ImageResponse, ImageObject +import time + +class MyCustomLLM(CustomLLM): + async def aimage_edit( + self, + model: str, + image: Any, + prompt: str, + model_response: ImageResponse, + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict, + logging_obj: Any, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[AsyncHTTPHandler] = None, + ) -> ImageResponse: + # Your custom image edit logic here + # e.g., call Stability AI, Black Forest Labs, etc. + return ImageResponse( + created=int(time.time()), + data=[ImageObject(url="https://example.com/edited-image.png")], + ) + +my_custom_llm = MyCustomLLM() +``` + + +2. Add to `config.yaml` + +In the config below, we pass + +python_filename: `custom_handler.py` +custom_handler_instance_name: `my_custom_llm`. This is defined in Step 1 + +custom_handler: `custom_handler.my_custom_llm` + +```yaml +model_list: + - model_name: "my-custom-image-edit-model" + litellm_params: + model: "my-custom-llm/my-model" + +litellm_settings: + custom_provider_map: + - {"provider": "my-custom-llm", "custom_handler": custom_handler.my_custom_llm} +``` + +```bash +litellm --config /path/to/config.yaml +``` + +3. Test it! + +```bash +curl -X POST 'http://0.0.0.0:4000/v1/images/edits' \ +-H 'Authorization: Bearer sk-1234' \ +-F 'model=my-custom-image-edit-model' \ +-F 'image=@/path/to/image.png' \ +-F 'prompt=Make the sky blue' +``` + +Expected Response + +``` +{ + "created": 1721955063, + "data": [{"url": "https://example.com/edited-image.png"}], +} +``` + ## Anthropic `/v1/messages` - Write the integration for .acompletion @@ -517,4 +595,34 @@ class CustomLLM(BaseLLM): client: Optional[AsyncHTTPHandler] = None, ) -> ImageResponse: raise CustomLLMError(status_code=500, message="Not implemented yet!") + + def image_edit( + self, + model: str, + image: Any, + prompt: str, + model_response: ImageResponse, + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict, + logging_obj: Any, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[HTTPHandler] = None, + ) -> ImageResponse: + raise CustomLLMError(status_code=500, message="Not implemented yet!") + + async def aimage_edit( + self, + model: str, + image: Any, + prompt: str, + model_response: ImageResponse, + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict, + logging_obj: Any, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[AsyncHTTPHandler] = None, + ) -> ImageResponse: + raise CustomLLMError(status_code=500, message="Not implemented yet!") ``` diff --git a/litellm/images/main.py b/litellm/images/main.py index 4aae96bf715..ca2d4e0b911 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -702,6 +702,59 @@ def image_edit( custom_llm_provider=custom_llm_provider, ) + # Check for custom provider + if custom_llm_provider in litellm._custom_providers: + custom_handler: Optional[CustomLLM] = None + for item in litellm.custom_provider_map: + if item["provider"] == custom_llm_provider: + custom_handler = item["custom_handler"] + + if custom_handler is None: + raise LiteLLMUnknownProvider( + model=model, custom_llm_provider=custom_llm_provider + ) + + model_response = ImageResponse() + + if _is_async: + async_custom_client: Optional[AsyncHTTPHandler] = None + if kwargs.get("client") is not None and isinstance( + kwargs.get("client"), AsyncHTTPHandler + ): + async_custom_client = kwargs.get("client") + + return custom_handler.aimage_edit( + model=model, + image=images, + prompt=prompt, + model_response=model_response, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + optional_params=kwargs, + logging_obj=litellm_logging_obj, + timeout=timeout, + client=async_custom_client, + ) + else: + custom_client: Optional[HTTPHandler] = None + if kwargs.get("client") is not None and isinstance( + kwargs.get("client"), HTTPHandler + ): + custom_client = kwargs.get("client") + + return custom_handler.image_edit( + model=model, + image=images, + prompt=prompt, + model_response=model_response, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + optional_params=kwargs, + logging_obj=litellm_logging_obj, + timeout=timeout, + client=custom_client, + ) + # get provider config image_edit_provider_config: Optional[BaseImageEditConfig] = ( ProviderConfigManager.get_provider_image_edit_config( diff --git a/litellm/llms/custom_llm.py b/litellm/llms/custom_llm.py index e88e8d5f1e3..d235df30f25 100644 --- a/litellm/llms/custom_llm.py +++ b/litellm/llms/custom_llm.py @@ -197,6 +197,36 @@ class CustomLLM(BaseLLM): ) -> EmbeddingResponse: raise CustomLLMError(status_code=500, message="Not implemented yet!") + def image_edit( + self, + model: str, + image: Any, + prompt: str, + model_response: ImageResponse, + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict, + logging_obj: Any, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[HTTPHandler] = None, + ) -> ImageResponse: + raise CustomLLMError(status_code=500, message="Not implemented yet!") + + async def aimage_edit( + self, + model: str, + image: Any, + prompt: str, + model_response: ImageResponse, + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict, + logging_obj: Any, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[AsyncHTTPHandler] = None, + ) -> ImageResponse: + raise CustomLLMError(status_code=500, message="Not implemented yet!") + def custom_chat_llm_router( async_fn: bool, stream: Optional[bool], custom_llm: CustomLLM diff --git a/tests/local_testing/test_custom_llm.py b/tests/local_testing/test_custom_llm.py index e61ede755e6..d0f32926551 100644 --- a/tests/local_testing/test_custom_llm.py +++ b/tests/local_testing/test_custom_llm.py @@ -309,6 +309,44 @@ class MyCustomLLM(CustomLLM): return model_response + def image_edit( + self, + model: str, + image: Any, + prompt: str, + model_response: ImageResponse, + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict, + logging_obj: Any, + timeout=None, + client: Optional[HTTPHandler] = None, + ) -> ImageResponse: + return ImageResponse( + created=int(time.time()), + data=[ImageObject(url="https://example.com/edited-image.png")], + response_ms=1000, + ) + + async def aimage_edit( + self, + model: str, + image: Any, + prompt: str, + model_response: ImageResponse, + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict, + logging_obj: Any, + timeout=None, + client: Optional[AsyncHTTPHandler] = None, + ) -> ImageResponse: + return ImageResponse( + created=int(time.time()), + data=[ImageObject(url="https://example.com/edited-image.png")], + response_ms=1000, + ) + def test_get_llm_provider(): """""" @@ -451,6 +489,69 @@ async def test_image_generation_async_additional_params(): } +def test_simple_image_edit(): + """Test sync image_edit with custom handler""" + my_custom_llm = MyCustomLLM() + litellm.custom_provider_map = [ + {"provider": "custom_llm", "custom_handler": my_custom_llm} + ] + resp = litellm.image_edit( + model="custom_llm/my-fake-model", + image=b"fake_image_bytes", + prompt="Edit this image", + ) + + print(resp) + assert resp.data[0].url == "https://example.com/edited-image.png" + + +@pytest.mark.asyncio +async def test_simple_image_edit_async(): + """Test async image_edit with custom handler""" + my_custom_llm = MyCustomLLM() + litellm.custom_provider_map = [ + {"provider": "custom_llm", "custom_handler": my_custom_llm} + ] + resp = await litellm.aimage_edit( + model="custom_llm/my-fake-model", + image=b"fake_image_bytes", + prompt="Edit this image", + ) + + print(resp) + assert resp.data[0].url == "https://example.com/edited-image.png" + + +@pytest.mark.asyncio +async def test_image_edit_async_additional_params(): + """Test that additional params are passed to custom handler""" + my_custom_llm = MyCustomLLM() + litellm.custom_provider_map = [ + {"provider": "custom_llm", "custom_handler": my_custom_llm} + ] + + with patch.object( + my_custom_llm, "aimage_edit", new=AsyncMock(return_value=ImageResponse( + created=int(time.time()), + data=[ImageObject(url="https://example.com/edited-image.png")], + )) + ) as mock_client: + resp = await litellm.aimage_edit( + model="custom_llm/my-fake-model", + image=b"fake_image_bytes", + prompt="Edit this image", + api_key="my-api-key", + api_base="my-api-base", + my_custom_param="my-custom-param", + ) + + print(resp) + + mock_client.assert_awaited_once() + assert mock_client.call_args.kwargs["api_key"] == "my-api-key" + assert mock_client.call_args.kwargs["api_base"] == "my-api-base" + + def test_get_supported_openai_params(): class MyCustomLLM(CustomLLM): From a4fb0df0281947d43dd1b3bceb9c2c69ae04a955 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 15 Dec 2025 17:40:58 -0800 Subject: [PATCH 012/121] [Feat] New provider - Agent Gateway, add pydantic ai agents (#18013) * init A2AProviderConfigManager * move file * move file * add pydnatic ai folder * init providers * test_pydantic_ai_non_streaming * fix import * INIT pydantic * use_a2a_form_fields * TestPydanticAITransformation --- .../litellm_completion_bridge/handler.py | 48 ++ litellm/a2a_protocol/providers/__init__.py | 11 + litellm/a2a_protocol/providers/base.py | 63 +++ .../a2a_protocol/providers/config_manager.py | 48 ++ .../providers/litellm_completion/README.md | 74 +++ .../providers/litellm_completion/__init__.py | 6 + .../providers/litellm_completion/handler.py | 295 ++++++++++ .../litellm_completion/transformation.py | 286 ++++++++++ .../providers/pydantic_ai_agents/__init__.py | 17 + .../providers/pydantic_ai_agents/config.py | 51 ++ .../providers/pydantic_ai_agents/handler.py | 106 ++++ .../pydantic_ai_agents/transformation.py | 523 ++++++++++++++++++ .../public_endpoints/agent_create_fields.json | 22 + tests/agent_tests/test_a2a.py | 163 +++++- .../test_pydantic_ai_agent_transformation.py | 99 ++++ .../public/assets/logos/pydantic.svg | 5 + .../src/components/agents/add_agent_form.tsx | 51 +- .../src/components/networking.tsx | 1 + 18 files changed, 1865 insertions(+), 4 deletions(-) create mode 100644 litellm/a2a_protocol/providers/__init__.py create mode 100644 litellm/a2a_protocol/providers/base.py create mode 100644 litellm/a2a_protocol/providers/config_manager.py create mode 100644 litellm/a2a_protocol/providers/litellm_completion/README.md create mode 100644 litellm/a2a_protocol/providers/litellm_completion/__init__.py create mode 100644 litellm/a2a_protocol/providers/litellm_completion/handler.py create mode 100644 litellm/a2a_protocol/providers/litellm_completion/transformation.py create mode 100644 litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py create mode 100644 litellm/a2a_protocol/providers/pydantic_ai_agents/config.py create mode 100644 litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py create mode 100644 litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py create mode 100644 tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py create mode 100644 ui/litellm-dashboard/public/assets/logos/pydantic.svg diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index cc222d3ee18..1916b04454a 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -18,6 +18,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( A2ACompletionBridgeTransformation, A2AStreamingContext, ) +from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager class A2ACompletionBridgeHandler: @@ -44,6 +45,29 @@ class A2ACompletionBridgeHandler: Returns: A2A SendMessageResponse dict """ + # Get provider config for custom_llm_provider + custom_llm_provider = litellm_params.get("custom_llm_provider") + a2a_provider_config = A2AProviderConfigManager.get_provider_config( + custom_llm_provider=custom_llm_provider + ) + + # If provider config exists, use it + if a2a_provider_config is not None: + if api_base is None: + raise ValueError(f"api_base is required for {custom_llm_provider}") + + verbose_logger.info( + f"A2A: Using provider config for {custom_llm_provider}" + ) + + response_data = await a2a_provider_config.handle_non_streaming( + request_id=request_id, + params=params, + api_base=api_base, + ) + + return response_data + # Extract message from params message = params.get("message", {}) @@ -119,6 +143,30 @@ class A2ACompletionBridgeHandler: Yields: A2A streaming response events """ + # Get provider config for custom_llm_provider + custom_llm_provider = litellm_params.get("custom_llm_provider") + a2a_provider_config = A2AProviderConfigManager.get_provider_config( + custom_llm_provider=custom_llm_provider + ) + + # If provider config exists, use it + if a2a_provider_config is not None: + if api_base is None: + raise ValueError(f"api_base is required for {custom_llm_provider}") + + verbose_logger.info( + f"A2A: Using provider config for {custom_llm_provider} (streaming)" + ) + + async for chunk in a2a_provider_config.handle_streaming( + request_id=request_id, + params=params, + api_base=api_base, + ): + yield chunk + + return + # Extract message from params message = params.get("message", {}) diff --git a/litellm/a2a_protocol/providers/__init__.py b/litellm/a2a_protocol/providers/__init__.py new file mode 100644 index 00000000000..873a5a83749 --- /dev/null +++ b/litellm/a2a_protocol/providers/__init__.py @@ -0,0 +1,11 @@ +""" +A2A Protocol Providers. + +This module contains provider-specific implementations for the A2A protocol. +""" + +from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig +from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager + +__all__ = ["BaseA2AProviderConfig", "A2AProviderConfigManager"] + diff --git a/litellm/a2a_protocol/providers/base.py b/litellm/a2a_protocol/providers/base.py new file mode 100644 index 00000000000..656bc78d6d2 --- /dev/null +++ b/litellm/a2a_protocol/providers/base.py @@ -0,0 +1,63 @@ +""" +Base configuration for A2A protocol providers. +""" + +from abc import ABC, abstractmethod +from typing import Any, AsyncIterator, Dict, Optional + + +class BaseA2AProviderConfig(ABC): + """ + Base configuration class for A2A protocol providers. + + Each provider should implement this interface to define how to handle + A2A requests for their specific agent type. + """ + + @abstractmethod + async def handle_non_streaming( + self, + request_id: str, + params: Dict[str, Any], + api_base: str, + **kwargs, + ) -> Dict[str, Any]: + """ + Handle non-streaming A2A request. + + Args: + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message + api_base: Base URL of the agent + **kwargs: Additional provider-specific parameters + + Returns: + A2A SendMessageResponse dict + """ + pass + + @abstractmethod + async def handle_streaming( + self, + request_id: str, + params: Dict[str, Any], + api_base: str, + **kwargs, + ) -> AsyncIterator[Dict[str, Any]]: + """ + Handle streaming A2A request. + + Args: + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message + api_base: Base URL of the agent + **kwargs: Additional provider-specific parameters + + Yields: + A2A streaming response events + """ + # This is an abstract method - subclasses must implement + # The yield is here to make this a generator function + if False: # pragma: no cover + yield {} + diff --git a/litellm/a2a_protocol/providers/config_manager.py b/litellm/a2a_protocol/providers/config_manager.py new file mode 100644 index 00000000000..e0703ec466b --- /dev/null +++ b/litellm/a2a_protocol/providers/config_manager.py @@ -0,0 +1,48 @@ +""" +A2A Provider Config Manager. + +Manages provider-specific configurations for A2A protocol. +""" + +from typing import Optional + +from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig + + +class A2AProviderConfigManager: + """ + Manager for A2A provider configurations. + + Similar to ProviderConfigManager in litellm.utils but specifically for A2A providers. + """ + + @staticmethod + def get_provider_config( + custom_llm_provider: Optional[str], + ) -> Optional[BaseA2AProviderConfig]: + """ + Get the provider configuration for a given custom_llm_provider. + + Args: + custom_llm_provider: The provider identifier (e.g., "pydantic_ai_agents") + + Returns: + Provider configuration instance or None if not found + """ + if custom_llm_provider is None: + return None + + if custom_llm_provider == "pydantic_ai_agents": + from litellm.a2a_protocol.providers.pydantic_ai_agents.config import ( + PydanticAIProviderConfig, + ) + + return PydanticAIProviderConfig() + + # Add more providers here as needed + # elif custom_llm_provider == "another_provider": + # from litellm.a2a_protocol.providers.another_provider.config import AnotherProviderConfig + # return AnotherProviderConfig() + + return None + diff --git a/litellm/a2a_protocol/providers/litellm_completion/README.md b/litellm/a2a_protocol/providers/litellm_completion/README.md new file mode 100644 index 00000000000..a809e9bf55e --- /dev/null +++ b/litellm/a2a_protocol/providers/litellm_completion/README.md @@ -0,0 +1,74 @@ +# A2A to LiteLLM Completion Bridge + +Routes A2A protocol requests through `litellm.acompletion`, enabling any LiteLLM-supported provider to be invoked via A2A. + +## Flow + +``` +A2A Request → Transform → litellm.acompletion → Transform → A2A Response +``` + +## SDK Usage + +Use the existing `asend_message` and `asend_message_streaming` functions with `litellm_params`: + +```python +from litellm.a2a_protocol import asend_message, asend_message_streaming +from a2a.types import SendMessageRequest, SendStreamingMessageRequest, MessageSendParams +from uuid import uuid4 + +# Non-streaming +request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={"role": "user", "parts": [{"kind": "text", "text": "Hello!"}], "messageId": uuid4().hex} + ) +) +response = await asend_message( + request=request, + api_base="http://localhost:2024", + litellm_params={"custom_llm_provider": "langgraph", "model": "agent"}, +) + +# Streaming +stream_request = SendStreamingMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={"role": "user", "parts": [{"kind": "text", "text": "Hello!"}], "messageId": uuid4().hex} + ) +) +async for chunk in asend_message_streaming( + request=stream_request, + api_base="http://localhost:2024", + litellm_params={"custom_llm_provider": "langgraph", "model": "agent"}, +): + print(chunk) +``` + +## Proxy Usage + +Configure an agent with `custom_llm_provider` in `litellm_params`: + +```yaml +agents: + - agent_name: my-langgraph-agent + agent_card_params: + name: "LangGraph Agent" + url: "http://localhost:2024" # Used as api_base + litellm_params: + custom_llm_provider: langgraph + model: agent +``` + +When an A2A request hits `/a2a/{agent_id}/message/send`, the bridge: + +1. Detects `custom_llm_provider` in agent's `litellm_params` +2. Transforms A2A message → OpenAI messages +3. Calls `litellm.acompletion(model="langgraph/agent", api_base="http://localhost:2024")` +4. Transforms response → A2A format + +## Classes + +- `A2ACompletionBridgeTransformation` - Static methods for message format conversion +- `A2ACompletionBridgeHandler` - Static methods for handling requests (streaming/non-streaming) + diff --git a/litellm/a2a_protocol/providers/litellm_completion/__init__.py b/litellm/a2a_protocol/providers/litellm_completion/__init__.py new file mode 100644 index 00000000000..3f2b88bfaa3 --- /dev/null +++ b/litellm/a2a_protocol/providers/litellm_completion/__init__.py @@ -0,0 +1,6 @@ +""" +LiteLLM Completion bridge provider for A2A protocol. + +Routes A2A requests through litellm.acompletion based on custom_llm_provider. +""" + diff --git a/litellm/a2a_protocol/providers/litellm_completion/handler.py b/litellm/a2a_protocol/providers/litellm_completion/handler.py new file mode 100644 index 00000000000..57388a5d0ed --- /dev/null +++ b/litellm/a2a_protocol/providers/litellm_completion/handler.py @@ -0,0 +1,295 @@ +""" +Handler for A2A to LiteLLM completion bridge. + +Routes A2A requests through litellm.acompletion based on custom_llm_provider. + +A2A Streaming Events (in order): +1. Task event (kind: "task") - Initial task creation with status "submitted" +2. Status update (kind: "status-update") - Status change to "working" +3. Artifact update (kind: "artifact-update") - Content/artifact delivery +4. Status update (kind: "status-update") - Final status "completed" with final=true +""" + +from typing import Any, AsyncIterator, Dict, Optional + +import litellm +from litellm._logging import verbose_logger +from litellm.a2a_protocol.litellm_completion_bridge.pydantic_ai_transformation import ( + PydanticAITransformation, +) +from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( + A2ACompletionBridgeTransformation, + A2AStreamingContext, +) + + +class A2ACompletionBridgeHandler: + """ + Static methods for handling A2A requests via LiteLLM completion. + """ + + @staticmethod + async def handle_non_streaming( + request_id: str, + params: Dict[str, Any], + litellm_params: Dict[str, Any], + api_base: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Handle non-streaming A2A request via litellm.acompletion. + + Args: + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message + litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.) + api_base: API base URL from agent_card_params + + Returns: + A2A SendMessageResponse dict + """ + # Check if this is a Pydantic AI agent request + custom_llm_provider = litellm_params.get("custom_llm_provider") + if custom_llm_provider == "pydantic_ai_agents": + if api_base is None: + raise ValueError("api_base is required for Pydantic AI agents") + + verbose_logger.info( + f"Pydantic AI: Routing to Pydantic AI agent at {api_base}" + ) + + # Send request directly to Pydantic AI agent + response_data = await PydanticAITransformation.send_non_streaming_request( + api_base=api_base, + request_id=request_id, + params=params, + ) + + return response_data + + # Extract message from params + message = params.get("message", {}) + + # Transform A2A message to OpenAI format + openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages( + message + ) + + # Get completion params + custom_llm_provider = litellm_params.get("custom_llm_provider") + model = litellm_params.get("model", "agent") + + # Build full model string if provider specified + # Skip prepending if model already starts with the provider prefix + if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"): + full_model = f"{custom_llm_provider}/{model}" + else: + full_model = model + + verbose_logger.info( + f"A2A completion bridge: model={full_model}, api_base={api_base}" + ) + + # Build completion params dict + completion_params = { + "model": full_model, + "messages": openai_messages, + "api_base": api_base, + "stream": False, + } + # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) + litellm_params_to_add = { + k: v for k, v in litellm_params.items() + if k not in ("model", "custom_llm_provider") + } + completion_params.update(litellm_params_to_add) + + # Call litellm.acompletion + response = await litellm.acompletion(**completion_params) + + # Transform response to A2A format + a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response( + response=response, + request_id=request_id, + ) + + verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}") + + return a2a_response + + @staticmethod + async def handle_streaming( + request_id: str, + params: Dict[str, Any], + litellm_params: Dict[str, Any], + api_base: Optional[str] = None, + ) -> AsyncIterator[Dict[str, Any]]: + """ + Handle streaming A2A request via litellm.acompletion with stream=True. + + Emits proper A2A streaming events: + 1. Task event (kind: "task") - Initial task with status "submitted" + 2. Status update (kind: "status-update") - Status "working" + 3. Artifact update (kind: "artifact-update") - Content delivery + 4. Status update (kind: "status-update") - Final "completed" status + + Args: + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message + litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.) + api_base: API base URL from agent_card_params + + Yields: + A2A streaming response events + """ + # Check if this is a Pydantic AI agent request + custom_llm_provider = litellm_params.get("custom_llm_provider") + if custom_llm_provider == "pydantic_ai_agents": + if api_base is None: + raise ValueError("api_base is required for Pydantic AI agents") + + verbose_logger.info( + f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}" + ) + + # Get non-streaming response first + response_data = await PydanticAITransformation.send_non_streaming_request( + api_base=api_base, + request_id=request_id, + params=params, + ) + + # Convert to fake streaming + async for chunk in PydanticAITransformation.fake_streaming_from_response( + response_data=response_data, + request_id=request_id, + ): + yield chunk + + return + + # Extract message from params + message = params.get("message", {}) + + # Create streaming context + ctx = A2AStreamingContext( + request_id=request_id, + input_message=message, + ) + + # Transform A2A message to OpenAI format + openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages( + message + ) + + # Get completion params + custom_llm_provider = litellm_params.get("custom_llm_provider") + model = litellm_params.get("model", "agent") + + # Build full model string if provider specified + # Skip prepending if model already starts with the provider prefix + if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"): + full_model = f"{custom_llm_provider}/{model}" + else: + full_model = model + + verbose_logger.info( + f"A2A completion bridge streaming: model={full_model}, api_base={api_base}" + ) + + # Build completion params dict + completion_params = { + "model": full_model, + "messages": openai_messages, + "api_base": api_base, + "stream": True, + } + # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) + litellm_params_to_add = { + k: v for k, v in litellm_params.items() + if k not in ("model", "custom_llm_provider") + } + completion_params.update(litellm_params_to_add) + + # 1. Emit initial task event (kind: "task", status: "submitted") + task_event = A2ACompletionBridgeTransformation.create_task_event(ctx) + yield task_event + + # 2. Emit status update (kind: "status-update", status: "working") + working_event = A2ACompletionBridgeTransformation.create_status_update_event( + ctx=ctx, + state="working", + final=False, + message_text="Processing request...", + ) + yield working_event + + # Call litellm.acompletion with streaming + response = await litellm.acompletion(**completion_params) + + # 3. Accumulate content and emit artifact update + accumulated_text = "" + chunk_count = 0 + async for chunk in response: # type: ignore[union-attr] + chunk_count += 1 + + # Extract delta content + content = "" + if chunk is not None and hasattr(chunk, "choices") and chunk.choices: + choice = chunk.choices[0] + if hasattr(choice, "delta") and choice.delta: + content = choice.delta.content or "" + + if content: + accumulated_text += content + + # Emit artifact update with accumulated content + if accumulated_text: + artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event( + ctx=ctx, + text=accumulated_text, + ) + yield artifact_event + + # 4. Emit final status update (kind: "status-update", status: "completed", final: true) + completed_event = A2ACompletionBridgeTransformation.create_status_update_event( + ctx=ctx, + state="completed", + final=True, + ) + yield completed_event + + verbose_logger.info( + f"A2A completion bridge streaming completed: request_id={request_id}, chunks={chunk_count}" + ) + + +# Convenience functions that delegate to the class methods +async def handle_a2a_completion( + request_id: str, + params: Dict[str, Any], + litellm_params: Dict[str, Any], + api_base: Optional[str] = None, +) -> Dict[str, Any]: + """Convenience function for non-streaming A2A completion.""" + return await A2ACompletionBridgeHandler.handle_non_streaming( + request_id=request_id, + params=params, + litellm_params=litellm_params, + api_base=api_base, + ) + + +async def handle_a2a_completion_streaming( + request_id: str, + params: Dict[str, Any], + litellm_params: Dict[str, Any], + api_base: Optional[str] = None, +) -> AsyncIterator[Dict[str, Any]]: + """Convenience function for streaming A2A completion.""" + async for chunk in A2ACompletionBridgeHandler.handle_streaming( + request_id=request_id, + params=params, + litellm_params=litellm_params, + api_base=api_base, + ): + yield chunk diff --git a/litellm/a2a_protocol/providers/litellm_completion/transformation.py b/litellm/a2a_protocol/providers/litellm_completion/transformation.py new file mode 100644 index 00000000000..bbe7daa9fc4 --- /dev/null +++ b/litellm/a2a_protocol/providers/litellm_completion/transformation.py @@ -0,0 +1,286 @@ +""" +Transformation utilities for A2A <-> OpenAI message format conversion. + +A2A Message Format: +{ + "role": "user", + "parts": [{"kind": "text", "text": "Hello!"}], + "messageId": "abc123" +} + +OpenAI Message Format: +{"role": "user", "content": "Hello!"} + +A2A Streaming Events: +- Task event (kind: "task") - Initial task creation with status "submitted" +- Status update (kind: "status-update") - Status changes (working, completed) +- Artifact update (kind: "artifact-update") - Content/artifact delivery +""" + +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional +from uuid import uuid4 + +from litellm._logging import verbose_logger + + +class A2AStreamingContext: + """ + Context holder for A2A streaming state. + Tracks task_id, context_id, and message accumulation. + """ + + def __init__(self, request_id: str, input_message: Dict[str, Any]): + self.request_id = request_id + self.task_id = str(uuid4()) + self.context_id = str(uuid4()) + self.input_message = input_message + self.accumulated_text = "" + self.has_emitted_task = False + self.has_emitted_working = False + + +class A2ACompletionBridgeTransformation: + """ + Static methods for transforming between A2A and OpenAI message formats. + """ + + @staticmethod + def a2a_message_to_openai_messages( + a2a_message: Dict[str, Any], + ) -> List[Dict[str, str]]: + """ + Transform an A2A message to OpenAI message format. + + Args: + a2a_message: A2A message with role, parts, and messageId + + Returns: + List of OpenAI-format messages + """ + role = a2a_message.get("role", "user") + parts = a2a_message.get("parts", []) + + # Map A2A roles to OpenAI roles + openai_role = role + if role == "user": + openai_role = "user" + elif role == "assistant": + openai_role = "assistant" + elif role == "system": + openai_role = "system" + + # Extract text content from parts + content_parts = [] + for part in parts: + kind = part.get("kind", "") + if kind == "text": + text = part.get("text", "") + content_parts.append(text) + + content = "\n".join(content_parts) if content_parts else "" + + verbose_logger.debug( + f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}" + ) + + return [{"role": openai_role, "content": content}] + + @staticmethod + def openai_response_to_a2a_response( + response: Any, + request_id: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Transform a LiteLLM ModelResponse to A2A SendMessageResponse format. + + Args: + response: LiteLLM ModelResponse object + request_id: Original A2A request ID + + Returns: + A2A SendMessageResponse dict + """ + # Extract content from response + content = "" + if hasattr(response, "choices") and response.choices: + choice = response.choices[0] + if hasattr(choice, "message") and choice.message: + content = choice.message.content or "" + + # Build A2A message + a2a_message = { + "role": "agent", + "parts": [{"kind": "text", "text": content}], + "messageId": uuid4().hex, + } + + # Build A2A response + a2a_response = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "message": a2a_message, + }, + } + + verbose_logger.debug( + f"OpenAI -> A2A transform: content_length={len(content)}" + ) + + return a2a_response + + @staticmethod + def _get_timestamp() -> str: + """Get current timestamp in ISO format with timezone.""" + return datetime.now(timezone.utc).isoformat() + + @staticmethod + def create_task_event( + ctx: A2AStreamingContext, + ) -> Dict[str, Any]: + """ + Create the initial task event with status 'submitted'. + + This is the first event emitted in an A2A streaming response. + """ + return { + "id": ctx.request_id, + "jsonrpc": "2.0", + "result": { + "contextId": ctx.context_id, + "history": [ + { + "contextId": ctx.context_id, + "kind": "message", + "messageId": ctx.input_message.get("messageId", uuid4().hex), + "parts": ctx.input_message.get("parts", []), + "role": ctx.input_message.get("role", "user"), + "taskId": ctx.task_id, + } + ], + "id": ctx.task_id, + "kind": "task", + "status": { + "state": "submitted", + }, + }, + } + + @staticmethod + def create_status_update_event( + ctx: A2AStreamingContext, + state: str, + final: bool = False, + message_text: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Create a status update event. + + Args: + ctx: Streaming context + state: Status state ('working', 'completed') + final: Whether this is the final event + message_text: Optional message text for 'working' status + """ + status: Dict[str, Any] = { + "state": state, + "timestamp": A2ACompletionBridgeTransformation._get_timestamp(), + } + + # Add message for 'working' status + if state == "working" and message_text: + status["message"] = { + "contextId": ctx.context_id, + "kind": "message", + "messageId": str(uuid4()), + "parts": [{"kind": "text", "text": message_text}], + "role": "agent", + "taskId": ctx.task_id, + } + + return { + "id": ctx.request_id, + "jsonrpc": "2.0", + "result": { + "contextId": ctx.context_id, + "final": final, + "kind": "status-update", + "status": status, + "taskId": ctx.task_id, + }, + } + + @staticmethod + def create_artifact_update_event( + ctx: A2AStreamingContext, + text: str, + ) -> Dict[str, Any]: + """ + Create an artifact update event with content. + + Args: + ctx: Streaming context + text: The text content for the artifact + """ + return { + "id": ctx.request_id, + "jsonrpc": "2.0", + "result": { + "artifact": { + "artifactId": str(uuid4()), + "name": "response", + "parts": [{"kind": "text", "text": text}], + }, + "contextId": ctx.context_id, + "kind": "artifact-update", + "taskId": ctx.task_id, + }, + } + + @staticmethod + def openai_chunk_to_a2a_chunk( + chunk: Any, + request_id: Optional[str] = None, + is_final: bool = False, + ) -> Optional[Dict[str, Any]]: + """ + Transform a LiteLLM streaming chunk to A2A streaming format. + + NOTE: This method is deprecated for streaming. Use the event-based + methods (create_task_event, create_status_update_event, + create_artifact_update_event) instead for proper A2A streaming. + + Args: + chunk: LiteLLM ModelResponse chunk + request_id: Original A2A request ID + is_final: Whether this is the final chunk + + Returns: + A2A streaming chunk dict or None if no content + """ + # Extract delta content + content = "" + if chunk is not None and hasattr(chunk, "choices") and chunk.choices: + choice = chunk.choices[0] + if hasattr(choice, "delta") and choice.delta: + content = choice.delta.content or "" + + if not content and not is_final: + return None + + # Build A2A streaming chunk (legacy format) + a2a_chunk = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "message": { + "role": "agent", + "parts": [{"kind": "text", "text": content}], + "messageId": uuid4().hex, + }, + "final": is_final, + }, + } + + return a2a_chunk diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py new file mode 100644 index 00000000000..2187400b2d1 --- /dev/null +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py @@ -0,0 +1,17 @@ +""" +Pydantic AI agent provider for A2A protocol. + +Pydantic AI agents follow A2A protocol but don't support streaming natively. +This provider handles fake streaming by converting non-streaming responses into streaming chunks. +""" + +from litellm.a2a_protocol.providers.pydantic_ai_agents.config import ( + PydanticAIProviderConfig, +) +from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAIHandler +from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( + PydanticAITransformation, +) + +__all__ = ["PydanticAIHandler", "PydanticAITransformation", "PydanticAIProviderConfig"] + diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py new file mode 100644 index 00000000000..acf09554e5e --- /dev/null +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py @@ -0,0 +1,51 @@ +""" +Pydantic AI provider configuration. +""" + +from typing import Any, AsyncIterator, Dict + +from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig +from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAIHandler + + +class PydanticAIProviderConfig(BaseA2AProviderConfig): + """ + Provider configuration for Pydantic AI agents. + + Pydantic AI agents follow A2A protocol but don't support streaming natively. + This config provides fake streaming by converting non-streaming responses into streaming chunks. + """ + + async def handle_non_streaming( + self, + request_id: str, + params: Dict[str, Any], + api_base: str, + **kwargs, + ) -> Dict[str, Any]: + """Handle non-streaming request to Pydantic AI agent.""" + return await PydanticAIHandler.handle_non_streaming( + request_id=request_id, + params=params, + api_base=api_base, + timeout=kwargs.get("timeout", 60.0), + ) + + async def handle_streaming( + self, + request_id: str, + params: Dict[str, Any], + api_base: str, + **kwargs, + ) -> AsyncIterator[Dict[str, Any]]: + """Handle streaming request with fake streaming.""" + async for chunk in PydanticAIHandler.handle_streaming( + request_id=request_id, + params=params, + api_base=api_base, + timeout=kwargs.get("timeout", 60.0), + chunk_size=kwargs.get("chunk_size", 50), + delay_ms=kwargs.get("delay_ms", 10), + ): + yield chunk + diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py new file mode 100644 index 00000000000..6680a9fe487 --- /dev/null +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py @@ -0,0 +1,106 @@ +""" +Handler for Pydantic AI agents. + +Pydantic AI agents follow A2A protocol but don't support streaming natively. +This handler provides fake streaming by converting non-streaming responses into streaming chunks. +""" + +from typing import Any, AsyncIterator, Dict + +from litellm._logging import verbose_logger +from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( + PydanticAITransformation, +) + + +class PydanticAIHandler: + """ + Handler for Pydantic AI agent requests. + + Provides: + - Direct non-streaming requests to Pydantic AI agents + - Fake streaming by converting non-streaming responses into streaming chunks + """ + + @staticmethod + async def handle_non_streaming( + request_id: str, + params: Dict[str, Any], + api_base: str, + timeout: float = 60.0, + ) -> Dict[str, Any]: + """ + Handle non-streaming request to Pydantic AI agent. + + Args: + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message + api_base: Base URL of the Pydantic AI agent + timeout: Request timeout in seconds + + Returns: + A2A SendMessageResponse dict + """ + verbose_logger.info( + f"Pydantic AI: Routing to Pydantic AI agent at {api_base}" + ) + + # Send request directly to Pydantic AI agent + response_data = await PydanticAITransformation.send_non_streaming_request( + api_base=api_base, + request_id=request_id, + params=params, + timeout=timeout, + ) + + return response_data + + @staticmethod + async def handle_streaming( + request_id: str, + params: Dict[str, Any], + api_base: str, + timeout: float = 60.0, + chunk_size: int = 50, + delay_ms: int = 10, + ) -> AsyncIterator[Dict[str, Any]]: + """ + Handle streaming request to Pydantic AI agent with fake streaming. + + Since Pydantic AI agents don't support streaming natively, this method: + 1. Makes a non-streaming request + 2. Converts the response into streaming chunks + + Args: + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message + api_base: Base URL of the Pydantic AI agent + timeout: Request timeout in seconds + chunk_size: Number of characters per chunk + delay_ms: Delay between chunks in milliseconds + + Yields: + A2A streaming response events + """ + verbose_logger.info( + f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}" + ) + + # Get raw task response first (not the transformed A2A format) + raw_response = await PydanticAITransformation.send_and_get_raw_response( + api_base=api_base, + request_id=request_id, + params=params, + timeout=timeout, + ) + + # Convert raw task response to fake streaming chunks + async for chunk in PydanticAITransformation.fake_streaming_from_response( + response_data=raw_response, + request_id=request_id, + chunk_size=chunk_size, + delay_ms=delay_ms, + ): + yield chunk + + diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py new file mode 100644 index 00000000000..67c2154f6c9 --- /dev/null +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -0,0 +1,523 @@ +""" +Transformation layer for Pydantic AI agents. + +Pydantic AI agents follow A2A protocol but don't support streaming. +This module provides fake streaming by converting non-streaming responses into streaming chunks. +""" + +import asyncio +from typing import Any, AsyncIterator, Dict, Optional +from uuid import uuid4 + +import httpx + +from litellm._logging import verbose_logger + + +class PydanticAITransformation: + """ + Transformation layer for Pydantic AI agents. + + Handles: + - Direct A2A requests to Pydantic AI endpoints + - Polling for task completion (since Pydantic AI doesn't support streaming) + - Fake streaming by chunking non-streaming responses + """ + + @staticmethod + def _remove_none_values(obj: Any) -> Any: + """ + Recursively remove None values from a dict/list structure. + + FastA2A/Pydantic AI servers don't accept None values for optional fields - + they expect those fields to be omitted entirely. + + Args: + obj: Dict, list, or other value to clean + + Returns: + 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 + } + elif isinstance(obj, list): + return [ + PydanticAITransformation._remove_none_values(item) + for item in obj + if item is not None + ] + else: + return obj + + @staticmethod + def _params_to_dict(params: Any) -> Dict[str, Any]: + """ + Convert params to a dict, handling Pydantic models. + + Args: + params: Dict or Pydantic model + + Returns: + Dict representation of params + """ + if hasattr(params, "model_dump"): + # Pydantic v2 model + return params.model_dump(mode="python", exclude_none=True) + elif hasattr(params, "dict"): + # Pydantic v1 model + return params.dict(exclude_none=True) + elif isinstance(params, dict): + return params + else: + # Try to convert to dict + return dict(params) + + @staticmethod + async def _poll_for_completion( + client: httpx.AsyncClient, + endpoint: str, + task_id: str, + request_id: str, + max_attempts: int = 30, + poll_interval: float = 0.5, + ) -> Dict[str, Any]: + """ + Poll for task completion using tasks/get method. + + Args: + client: HTTPX async client + endpoint: API endpoint URL + task_id: Task ID to poll for + request_id: JSON-RPC request ID + max_attempts: Maximum polling attempts + poll_interval: Seconds between poll attempts + + Returns: + Completed task response + """ + for attempt in range(max_attempts): + poll_request = { + "jsonrpc": "2.0", + "id": f"{request_id}-poll-{attempt}", + "method": "tasks/get", + "params": {"id": task_id}, + } + + response = await client.post( + endpoint, + json=poll_request, + headers={"Content-Type": "application/json"}, + ) + response.raise_for_status() + poll_data = response.json() + + result = poll_data.get("result", {}) + status = result.get("status", {}) + state = status.get("state", "") + + verbose_logger.debug( + f"Pydantic AI: Poll attempt {attempt + 1}/{max_attempts}, state={state}" + ) + + if state == "completed": + return poll_data + elif state in ("failed", "canceled"): + raise Exception(f"Task {task_id} ended with state: {state}") + + await asyncio.sleep(poll_interval) + + raise TimeoutError(f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds") + + @staticmethod + async def _send_and_poll_raw( + api_base: str, + request_id: str, + params: Any, + timeout: float = 60.0, + ) -> Dict[str, Any]: + """ + Send a request to Pydantic AI agent and return the raw task response. + + This is an internal method used by both non-streaming and streaming handlers. + Returns the raw Pydantic AI task format with history/artifacts. + + Args: + api_base: Base URL of the Pydantic AI agent + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message + timeout: Request timeout in seconds + + Returns: + 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) + + # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI + if "message" in params_dict: + params_dict["message"]["kind"] = "message" + + # Build A2A JSON-RPC request using message/send method for FastA2A compatibility + a2a_request = { + "jsonrpc": "2.0", + "id": request_id, + "method": "message/send", + "params": params_dict, + } + + # FastA2A uses root endpoint (/) not /messages + endpoint = api_base.rstrip("/") + + verbose_logger.info( + f"Pydantic AI: Sending non-streaming request to {endpoint}" + ) + + # Send request to Pydantic AI agent + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.post( + endpoint, + json=a2a_request, + headers={"Content-Type": "application/json"}, + ) + response.raise_for_status() + response_data = response.json() + + # Check if task is already completed + result = response_data.get("result", {}) + status = result.get("status", {}) + state = status.get("state", "") + + if state != "completed": + # Need to poll for completion + task_id = result.get("id") + if task_id: + verbose_logger.info( + f"Pydantic AI: Task {task_id} submitted, polling for completion..." + ) + response_data = await PydanticAITransformation._poll_for_completion( + client=client, + endpoint=endpoint, + task_id=task_id, + request_id=request_id, + ) + + verbose_logger.info(f"Pydantic AI: Received completed response for request_id={request_id}") + + return response_data + + @staticmethod + async def send_non_streaming_request( + api_base: str, + request_id: str, + params: Any, + timeout: float = 60.0, + ) -> Dict[str, Any]: + """ + Send a non-streaming A2A request to Pydantic AI agent and wait for completion. + + Args: + api_base: Base URL of the Pydantic AI agent (e.g., "http://localhost:9999") + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message (dict or Pydantic model) + timeout: Request timeout in seconds + + Returns: + Standard A2A non-streaming response format with message + """ + # Get raw task response + raw_response = await PydanticAITransformation._send_and_poll_raw( + api_base=api_base, + request_id=request_id, + params=params, + timeout=timeout, + ) + + # Transform to standard A2A non-streaming format + return PydanticAITransformation._transform_to_a2a_response( + response_data=raw_response, + request_id=request_id, + ) + + @staticmethod + async def send_and_get_raw_response( + api_base: str, + request_id: str, + params: Any, + timeout: float = 60.0, + ) -> Dict[str, Any]: + """ + Send a request to Pydantic AI agent and return the raw task response. + + Used by streaming handler to get raw response for fake streaming. + + Args: + api_base: Base URL of the Pydantic AI agent + request_id: A2A JSON-RPC request ID + params: A2A MessageSendParams containing the message + timeout: Request timeout in seconds + + Returns: + Raw Pydantic AI task response (with history/artifacts) + """ + return await PydanticAITransformation._send_and_poll_raw( + api_base=api_base, + request_id=request_id, + params=params, + timeout=timeout, + ) + + @staticmethod + def _transform_to_a2a_response( + response_data: Dict[str, Any], + request_id: str, + ) -> Dict[str, Any]: + """ + Transform Pydantic AI task response to standard A2A non-streaming format. + + Pydantic AI returns a task with history/artifacts, but the standard A2A + non-streaming format expects: + { + "jsonrpc": "2.0", + "id": "...", + "result": { + "message": { + "role": "agent", + "parts": [{"kind": "text", "text": "..."}], + "messageId": "..." + } + } + } + + Args: + response_data: Pydantic AI task response + request_id: Original request ID + + Returns: + Standard A2A non-streaming response format + """ + # Extract the agent response text + full_text, message_id, parts = PydanticAITransformation._extract_response_text( + response_data + ) + + # Build standard A2A message + a2a_message = { + "role": "agent", + "parts": parts if parts else [{"kind": "text", "text": full_text}], + "messageId": message_id, + } + + # Return standard A2A non-streaming format + return { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "message": a2a_message, + }, + } + + @staticmethod + def _extract_response_text(response_data: Dict[str, Any]) -> tuple[str, str, list]: + """ + Extract response text from completed task response. + + Pydantic AI returns completed tasks with: + - history: list of messages (user and agent) + - artifacts: list of result artifacts + + Args: + response_data: Completed task response + + Returns: + Tuple of (full_text, message_id, parts) + """ + result = response_data.get("result", {}) + + # Try to extract from artifacts first (preferred for results) + artifacts = result.get("artifacts", []) + if artifacts: + for artifact in artifacts: + parts = artifact.get("parts", []) + for part in parts: + if part.get("kind") == "text": + text = part.get("text", "") + if text: + return text, str(uuid4()), parts + + # Fall back to history - get the last agent message + history = result.get("history", []) + for msg in reversed(history): + if msg.get("role") == "agent": + parts = msg.get("parts", []) + message_id = msg.get("messageId", str(uuid4())) + full_text = "" + for part in parts: + if part.get("kind") == "text": + full_text += part.get("text", "") + if full_text: + return full_text, message_id, parts + + # Fall back to message field (original format) + message = result.get("message", {}) + if message: + parts = message.get("parts", []) + message_id = message.get("messageId", str(uuid4())) + full_text = "" + for part in parts: + if part.get("kind") == "text": + full_text += part.get("text", "") + return full_text, message_id, parts + + return "", str(uuid4()), [] + + @staticmethod + async def fake_streaming_from_response( + response_data: Dict[str, Any], + request_id: str, + chunk_size: int = 50, + delay_ms: int = 10, + ) -> AsyncIterator[Dict[str, Any]]: + """ + Convert a non-streaming A2A response into fake streaming chunks. + + Emits proper A2A streaming events: + 1. Task event (kind: "task") - Initial task with status "submitted" + 2. Status update (kind: "status-update") - Status "working" + 3. Artifact update chunks (kind: "artifact-update") - Content delivery in chunks + 4. Status update (kind: "status-update") - Final "completed" status + + Args: + response_data: Non-streaming A2A response dict (completed task) + request_id: A2A JSON-RPC request ID + chunk_size: Number of characters per chunk (default: 50) + delay_ms: Delay between chunks in milliseconds (default: 10) + + Yields: + A2A streaming response events + """ + # Extract the response text from completed task + full_text, message_id, parts = PydanticAITransformation._extract_response_text( + response_data + ) + + # Extract input message from raw response for history + result = response_data.get("result", {}) + history = result.get("history", []) + input_message = {} + for msg in history: + if msg.get("role") == "user": + input_message = msg + break + + # Generate IDs for streaming events + task_id = str(uuid4()) + context_id = str(uuid4()) + artifact_id = str(uuid4()) + input_message_id = input_message.get("messageId", str(uuid4())) + + # 1. Emit initial task event (kind: "task", status: "submitted") + # Format matches A2ACompletionBridgeTransformation.create_task_event + task_event = { + "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 = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": False, + "kind": "status-update", + "status": { + "state": "working", + }, + "taskId": task_id, + }, + } + yield working_event + + # Small delay to simulate processing + await asyncio.sleep(delay_ms / 1000.0) + + # 3. Emit artifact update chunks (kind: "artifact-update") + # Format matches A2ACompletionBridgeTransformation.create_artifact_update_event + if 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) + + 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, + } + ], + }, + }, + } + yield artifact_event + + # Add delay between chunks (except for last chunk) + if not is_last_chunk: + await asyncio.sleep(delay_ms / 1000.0) + + # 4. Emit final status update (kind: "status-update", status: "completed", final: true) + completed_event = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": True, + "kind": "status-update", + "status": { + "state": "completed", + }, + "taskId": task_id, + }, + } + yield completed_event + + verbose_logger.info( + f"Pydantic AI: Fake streaming completed for request_id={request_id}" + ) + + diff --git a/litellm/proxy/public_endpoints/agent_create_fields.json b/litellm/proxy/public_endpoints/agent_create_fields.json index 347a58a7675..c559b61a76f 100644 --- a/litellm/proxy/public_endpoints/agent_create_fields.json +++ b/litellm/proxy/public_endpoints/agent_create_fields.json @@ -144,6 +144,28 @@ "litellm_params_template": { "custom_llm_provider": "azure_ai" } + }, + { + "agent_type": "pydantic_ai_agents", + "agent_type_display_name": "Pydantic AI", + "description": "Connect to Pydantic AI agents via A2A protocol (with fake streaming support)", + "logo_url": "/ui/assets/logos/pydantic.svg", + "use_a2a_form_fields": true, + "credential_fields": [ + { + "key": "api_base", + "label": "Agent URL", + "placeholder": "http://localhost:9999", + "tooltip": "The base URL for your Pydantic AI agent server", + "required": true, + "field_type": "text", + "default_value": "http://localhost:9999", + "include_in_litellm_params": true + } + ], + "litellm_params_template": { + "custom_llm_provider": "pydantic_ai_agents" + } } ] diff --git a/tests/agent_tests/test_a2a.py b/tests/agent_tests/test_a2a.py index eeab2680564..1550d61f7b0 100644 --- a/tests/agent_tests/test_a2a.py +++ b/tests/agent_tests/test_a2a.py @@ -21,10 +21,7 @@ from litellm.types.utils import StandardLoggingPayload sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path - from a2a.types import MessageSendParams, SendMessageRequest - - @pytest.mark.asyncio async def test_asend_message_with_client_decorator(): """ @@ -165,3 +162,163 @@ async def test_a2a_logging_payload(): # This confirms the A2A cost calculator is working assert response_cost is not None, "response_cost should not be None" assert response_cost == 0.0, f"response_cost should be 0.0 for A2A, got: {response_cost}" + + +@pytest.mark.asyncio +async def test_pydantic_ai_non_streaming(): + """ + Test non-streaming requests to Pydantic AI agents. + + Pydantic AI agents follow A2A protocol but don't support streaming. + This test validates non-streaming requests work correctly. + """ + litellm._turn_on_debug() + from litellm.a2a_protocol import asend_message + + # Build the request + send_message_payload = { + "message": { + "role": "user", + "parts": [ + { + "kind": "text", + "text": "Hello from Pydantic AI test!", + } + ], + "messageId": uuid4().hex, + }, + } + + request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams(**send_message_payload), + ) + + # Send message using Pydantic AI provider + response = await asend_message( + request=request, + api_base="http://localhost:9999", + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + + # Print response for debugging + print("\n=== Pydantic AI Non-Streaming Response ===") + print(response.model_dump(mode="json", exclude_none=True)) + + # Basic assertions + assert response is not None + assert hasattr(response, "result") + + # Verify result structure + result = response.result + assert result is not None + + # Pydantic AI returns a task with history/artifacts, not a direct message + # Check for either format + result_dict = result if isinstance(result, dict) else result.model_dump(mode="python", exclude_none=True) + has_message = "message" in result_dict + has_history = "history" in result_dict + has_artifacts = "artifacts" in result_dict + + assert has_message or has_history or has_artifacts, ( + f"Result should contain 'message', 'history', or 'artifacts'. Got: {list(result_dict.keys())}" + ) + + # If it's a task response (Pydantic AI style), verify we got agent response + if has_history: + history = result_dict.get("history", []) + agent_messages = [m for m in history if m.get("role") == "agent"] + assert len(agent_messages) > 0, "Should have at least one agent message in history" + + # Verify agent message has text content + agent_msg = agent_messages[-1] + parts = agent_msg.get("parts", []) + text_parts = [p for p in parts if p.get("kind") == "text"] + assert len(text_parts) > 0, "Agent message should have text content" + print(f"\nAgent response: {text_parts[0].get('text')}") + + +@pytest.mark.asyncio +async def test_pydantic_ai_fake_streaming(): + """ + Test fake streaming for Pydantic AI agents. + + Pydantic AI agents don't support streaming natively. + This test validates that fake streaming works by converting + non-streaming responses into streaming chunks. + """ + litellm._turn_on_debug() + from litellm.a2a_protocol import asend_message_streaming + + # Build the request + from a2a.types import SendStreamingMessageRequest + + send_message_payload = { + "message": { + "role": "user", + "parts": [ + { + "kind": "text", + "text": "Hello from Pydantic AI streaming test!", + } + ], + "messageId": uuid4().hex, + }, + } + + request = SendStreamingMessageRequest( + id=str(uuid4()), + params=MessageSendParams(**send_message_payload), + ) + + # Send streaming message using Pydantic AI provider + print("\n=== Pydantic AI Fake Streaming Response ===") + chunks_received = 0 + task_event_received = False + working_event_received = False + artifact_event_received = False + completed_event_received = False + + async for chunk in asend_message_streaming( + request=request, + api_base="http://localhost:9999", + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ): + chunks_received += 1 + print(f"\nChunk {chunks_received}:") + + # Convert chunk to dict for inspection + chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else chunk + print(json.dumps(chunk_dict, indent=2)) + + # Check event types + result = chunk_dict.get("result", {}) + kind = result.get("kind") + + if kind == "task": + task_event_received = True + elif kind == "status-update": + status = result.get("status", {}) + state = status.get("state") + if state == "working": + working_event_received = True + elif state == "completed": + completed_event_received = True + elif kind == "artifact-update": + artifact_event_received = True + + print(f"\n=== Streaming Summary ===") + print(f"Total chunks received: {chunks_received}") + print(f"Task event received: {task_event_received}") + print(f"Working event received: {working_event_received}") + print(f"Artifact event received: {artifact_event_received}") + print(f"Completed event received: {completed_event_received}") + + # Verify we received chunks + assert chunks_received > 0, "Should receive at least one chunk" + + # Verify all required event types were received + assert task_event_received, "Should receive task event" + assert working_event_received, "Should receive working status event" + assert artifact_event_received, "Should receive artifact update event" + assert completed_event_received, "Should receive completed status event" diff --git a/tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py b/tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py new file mode 100644 index 00000000000..7b10c46c2fd --- /dev/null +++ b/tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py @@ -0,0 +1,99 @@ +""" +Tests for Pydantic AI agents transformation. + +Tests the helper functions and response transformation without making real API calls. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( + PydanticAITransformation, +) + + +class TestPydanticAITransformation: + """Tests for PydanticAITransformation helper methods.""" + + def test_remove_none_values(self): + """ + Test that _remove_none_values recursively removes None values from dicts. + FastA2A servers reject None values for optional fields. + """ + input_data = { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "contextId": None, + "taskId": None, + "metadata": None, + }, + "configuration": None, + "metadata": {"key": "value", "empty": None}, + } + + result = PydanticAITransformation._remove_none_values(input_data) + + # None values should be removed + assert "contextId" not in result["message"] + assert "taskId" not in result["message"] + assert "metadata" not in result["message"] + assert "configuration" not in result + assert "empty" not in result["metadata"] + + # Non-None values should be preserved + assert result["message"]["role"] == "user" + assert result["message"]["parts"] == [{"kind": "text", "text": "Hello"}] + assert result["metadata"]["key"] == "value" + + def test_transform_to_a2a_response(self): + """ + Test that _transform_to_a2a_response converts Pydantic AI task format + to standard A2A non-streaming response format. + """ + # Pydantic AI returns tasks with history/artifacts + pydantic_ai_response = { + "jsonrpc": "2.0", + "id": "req-123", + "result": { + "id": "task-456", + "kind": "task", + "status": {"state": "completed"}, + "history": [ + { + "role": "user", + "parts": [{"kind": "text", "text": "What is 2+2?"}], + "messageId": "msg-user-1", + }, + { + "role": "agent", + "parts": [{"kind": "text", "text": "The answer is 4."}], + "messageId": "msg-agent-1", + }, + ], + "artifacts": [ + { + "artifactId": "artifact-1", + "name": "response", + "parts": [{"kind": "text", "text": "The answer is 4."}], + } + ], + }, + } + + result = PydanticAITransformation._transform_to_a2a_response( + response_data=pydantic_ai_response, + request_id="req-123", + ) + + # Should return standard A2A format with message + assert result["jsonrpc"] == "2.0" + assert result["id"] == "req-123" + assert "message" in result["result"] + assert result["result"]["message"]["role"] == "agent" + assert result["result"]["message"]["parts"][0]["text"] == "The answer is 4." + diff --git a/ui/litellm-dashboard/public/assets/logos/pydantic.svg b/ui/litellm-dashboard/public/assets/logos/pydantic.svg new file mode 100644 index 00000000000..0ff8e5c44c7 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/pydantic.svg @@ -0,0 +1,5 @@ + + + diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx index f44d71cbada..f4e0137bd06 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, message, Select } from "antd"; +import { Modal, Form, message, Select, Input } from "antd"; import { Button } from "@tremor/react"; import { createAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking"; import AgentFormFields from "./agent_form_fields"; @@ -57,6 +57,26 @@ const AddAgentForm: React.FC = ({ if (agentType === "a2a") { agentData = buildAgentDataFromForm(values); + } else if (selectedAgentTypeInfo?.use_a2a_form_fields) { + // A2A-compatible agents use the standard A2A form builder + // but need to add litellm_params from the agent type config + agentData = buildAgentDataFromForm(values); + + // Merge litellm_params_template + if (selectedAgentTypeInfo.litellm_params_template) { + agentData.litellm_params = { + ...agentData.litellm_params, + ...selectedAgentTypeInfo.litellm_params_template, + }; + } + + // Add credential fields to litellm_params + for (const field of selectedAgentTypeInfo.credential_fields) { + const value = values[field.key]; + if (value && field.include_in_litellm_params !== false) { + agentData.litellm_params[field.key] = value; + } + } } else if (selectedAgentTypeInfo) { agentData = buildDynamicAgentData(values, selectedAgentTypeInfo); } @@ -167,6 +187,35 @@ const AddAgentForm: React.FC = ({
{agentType === "a2a" ? ( + ) : selectedAgentTypeInfo?.use_a2a_form_fields ? ( + // A2A-compatible agents (like Pydantic AI) use full A2A form fields + // plus any additional credential fields + <> + + {selectedAgentTypeInfo.credential_fields.length > 0 && ( +
+

+ {selectedAgentTypeInfo.agent_type_display_name} Settings +

+ {selectedAgentTypeInfo.credential_fields.map((field) => ( + + {field.field_type === "password" ? ( + + ) : ( + + )} + + ))} +
+ )} + ) : selectedAgentTypeInfo ? ( ) : null} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 65c41c5aab4..ddabc6f5212 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -215,6 +215,7 @@ export interface AgentCreateInfo { credential_fields: AgentCredentialFieldMetadata[]; litellm_params_template?: Record | null; model_template?: string | null; + use_a2a_form_fields?: boolean; } export interface PublicModelHubInfo { From 4fdbbdfe6de73637e631ff41ad74da8f6e104eab Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Mon, 15 Dec 2025 22:57:40 -0300 Subject: [PATCH 013/121] fix(anthropic): correct claude-3-7-sonnet max_tokens to 64K default (#17979) Claude 3.7 Sonnet's default max_output_tokens is 64000, not 128000. The 128K output limit requires the beta header 'output-128k-2025-02-19'. This fixes the integration test failure where requests with max_tokens=128000 were being rejected by the Anthropic API. Fixes test_multiturn_tool_calls in test_anthropic_responses_api.py --- litellm/model_prices_and_context_window_backup.json | 8 ++++---- model_prices_and_context_window.json | 8 ++++---- .../chat/test_anthropic_chat_transformation.py | 11 ++++++----- 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4b016bc6ca6..26aae425fae 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6723,8 +6723,8 @@ "input_cost_per_token": 3e-06, "litellm_provider": "anthropic", "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { @@ -6752,8 +6752,8 @@ "input_cost_per_token": 3e-06, "litellm_provider": "anthropic", "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4b016bc6ca6..26aae425fae 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6723,8 +6723,8 @@ "input_cost_per_token": 3e-06, "litellm_provider": "anthropic", "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { @@ -6752,8 +6752,8 @@ "input_cost_per_token": 3e-06, "litellm_provider": "anthropic", "max_input_tokens": 200000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index ec612109d9c..9b6d1c6e178 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1651,15 +1651,16 @@ def test_get_max_tokens_for_model_claude_35(): def test_get_max_tokens_for_model_claude_37(): """ Test that get_max_tokens_for_model returns correct value for Claude 3.7 models. - Claude 3.7 Sonnet has max_output_tokens of 128000 (128K with extended thinking). + Claude 3.7 Sonnet has max_output_tokens of 64000 by default. + 128K output requires the beta header 'output-128k-2025-02-19'. Fixes: https://github.com/BerriAI/litellm/issues/8835 """ config = AnthropicConfig() - # Claude 3.7 Sonnet should return 128000 (128K) + # Claude 3.7 Sonnet should return 64000 (64K default, 128K requires beta header) max_tokens = config.get_max_tokens_for_model("claude-3-7-sonnet-20250219") - assert max_tokens == 128000 + assert max_tokens == 64000 def test_get_max_tokens_for_model_unknown(): @@ -1698,9 +1699,9 @@ def test_get_config_with_model_uses_dynamic_max_tokens(): config_claude35 = AnthropicConfig.get_config(model="claude-3-5-sonnet-20241022") assert config_claude35["max_tokens"] == 8192 - # Claude 3.7 model should get 128000 (128K with extended thinking) + # Claude 3.7 model should get 64000 (64K default, 128K requires beta header) config_claude37 = AnthropicConfig.get_config(model="claude-3-7-sonnet-20250219") - assert config_claude37["max_tokens"] == 128000 + assert config_claude37["max_tokens"] == 64000 def test_get_config_without_model_uses_fallback(): From b0a9c85a522bdd3eaa1e6b8fabbc37c0228b4c80 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 16 Dec 2025 07:38:07 +0530 Subject: [PATCH 014/121] fix: fix ruff linting errors --- litellm/a2a_protocol/providers/base.py | 2 +- .../a2a_protocol/providers/pydantic_ai_agents/transformation.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/a2a_protocol/providers/base.py b/litellm/a2a_protocol/providers/base.py index 656bc78d6d2..9931076a948 100644 --- a/litellm/a2a_protocol/providers/base.py +++ b/litellm/a2a_protocol/providers/base.py @@ -3,7 +3,7 @@ Base configuration for A2A protocol providers. """ from abc import ABC, abstractmethod -from typing import Any, AsyncIterator, Dict, Optional +from typing import Any, AsyncIterator, Dict class BaseA2AProviderConfig(ABC): diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 67c2154f6c9..6f46933cf9f 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,7 +6,7 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from typing import Any, AsyncIterator, Dict, Optional +from typing import Any, AsyncIterator, Dict from uuid import uuid4 import httpx From c754794bc3171d64a52dda8d2f82af963a9432da Mon Sep 17 00:00:00 2001 From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Date: Mon, 15 Dec 2025 18:08:21 -0800 Subject: [PATCH 015/121] [fix] add qwen3-embedding-8b input per token price (#18018) * added embedding input token price * added embedding input token price --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 26aae425fae..01d7f076edc 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -30632,7 +30632,7 @@ "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "input_cost_per_token": 0.0, + "input_cost_per_token": 1e-07, "output_cost_per_token": 0.0, "litellm_provider": "fireworks_ai", "mode": "embedding" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 26aae425fae..2a7f8aa3ddf 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -30628,11 +30628,11 @@ "litellm_provider": "fireworks_ai", "mode": "embedding" }, - "fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": { + "fireworks_ai/accounts/fireworks/models/": { "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "input_cost_per_token": 0.0, + "input_cost_per_token": 1e-07, "output_cost_per_token": 0.0, "litellm_provider": "fireworks_ai", "mode": "embedding" From 999ffabc39cd5175293245160d9a4e8e8283fd61 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Mon, 15 Dec 2025 23:11:02 -0300 Subject: [PATCH 016/121] fix(gemini): use JSON instead of form-data for image edit requests (#18012) * fix(gemini): use JSON instead of form-data for image edit requests Gemini's image edit API expects JSON body, not multipart/form-data. The handler was sending form-encoded data which caused 400 errors: "Invalid JSON payload received. Unexpected token." Changes: - Add use_multipart_form_data() method to BaseImageEditConfig (default True) - Modify image_edit_handler to use json= when use_multipart_form_data() is False - Override use_multipart_form_data() in GeminiImageEditConfig to return False * test(gemini): add test for use_multipart_form_data --- .../base_llm/image_edit/transformation.py | 9 ++++ litellm/llms/custom_httpx/llm_http_handler.py | 50 +++++++++++++------ .../llms/gemini/image_edit/transformation.py | 4 ++ .../test_gemini_image_edit_transformation.py | 11 ++++ 4 files changed, 60 insertions(+), 14 deletions(-) diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index f3ae2d32eaa..d522675296f 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -109,6 +109,15 @@ class BaseImageEditConfig(ABC): ) -> ImageResponse: pass + def use_multipart_form_data(self) -> bool: + """ + Return True if the provider uses multipart/form-data for image edit requests. + Return False if the provider uses JSON requests. + + Default is True for backwards compatibility with OpenAI-style providers. + """ + return True + def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a7789a181f..4b38f542159 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -3768,13 +3768,24 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=api_base, - headers=headers, - data=data, - files=files, - timeout=timeout, - ) + # Check if provider uses multipart/form-data or JSON + if image_edit_provider_config.use_multipart_form_data(): + # Use form-data (OpenAI style) + response = sync_httpx_client.post( + url=api_base, + headers=headers, + data=data, + files=files, + timeout=timeout, + ) + else: + # Use JSON (Gemini style) + response = sync_httpx_client.post( + url=api_base, + headers=headers, + json=data, + timeout=timeout, + ) except Exception as e: raise self._handle_error( @@ -3853,13 +3864,24 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=api_base, - headers=headers, - data=data, - files=files, - timeout=timeout, - ) + # Check if provider uses multipart/form-data or JSON + if image_edit_provider_config.use_multipart_form_data(): + # Use form-data (OpenAI style) + response = await async_httpx_client.post( + url=api_base, + headers=headers, + data=data, + files=files, + timeout=timeout, + ) + else: + # Use JSON (Gemini style) + response = await async_httpx_client.post( + url=api_base, + headers=headers, + json=data, + timeout=timeout, + ) except Exception as e: raise self._handle_error( diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 830c58a0062..78a7ff9546f 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -63,6 +63,10 @@ class GeminiImageEditConfig(BaseImageEditConfig): headers["Content-Type"] = "application/json" return headers + def use_multipart_form_data(self) -> bool: + """Gemini uses JSON requests, not multipart/form-data.""" + return False + def get_complete_url( self, model: str, diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index 2732bf1595a..021cfaeff5e 100644 --- a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -147,3 +147,14 @@ class TestGeminiImageEditTransformation: headers={}, ) + def test_use_multipart_form_data_returns_false(self) -> None: + """ + Gemini uses JSON requests, not multipart/form-data. + This is critical because httpx sends data differently: + - data=dict sends form-encoded + - json=dict sends JSON + + Without this, Gemini returns: "Invalid JSON payload received. Unexpected token." + """ + assert self.config.use_multipart_form_data() is False + From fabbd28b4578de139e1dceae8d514eccfd4ae0aa Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 15 Dec 2025 18:31:35 -0800 Subject: [PATCH 017/121] Adding Milvus to Vector Store in UI --- .../src/components/vector_store_providers.tsx | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index a6f222da67c..68c5c3745c3 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -4,6 +4,7 @@ export enum VectorStoreProviders { VertexRagEngine = "Vertex AI RAG Engine", OpenAI = "OpenAI", Azure = "Azure OpenAI", + Milvus = "Milvus", } export const vectorStoreProviderMap: Record = { @@ -12,6 +13,7 @@ export const vectorStoreProviderMap: Record = { VertexRagEngine: "vertex_ai", OpenAI: "openai", Azure: "azure", + Milvus: "milvus", }; const asset_logos_folder = "../ui/assets/logos/"; @@ -22,6 +24,7 @@ export const vectorStoreProviderLogoMap: Record = { [VectorStoreProviders.VertexRagEngine]: `${asset_logos_folder}google.svg`, [VectorStoreProviders.OpenAI]: `${asset_logos_folder}openai_small.svg`, [VectorStoreProviders.Azure]: `${asset_logos_folder}microsoft_azure.svg`, + [VectorStoreProviders.Milvus]: `${asset_logos_folder}milvus.svg`, }; // Define field types for provider-specific configurations @@ -84,6 +87,25 @@ export const vectorStoreProviderFields: Record type: "text", }, ], + milvus: [ + { + name: "api_key", + label: "API Key", + tooltip: + "To obtain a token, you should use a colon (:) to concatenate the username and password that you use to access your Milvus instance (e.g., username:password)", + placeholder: "username:password or api key", + required: true, + type: "password", + }, + { + name: "api_base", + label: "API Base", + tooltip: "Enter your Milvus endpoint (e.g., https://your-milvus-endpoint.com/)", + placeholder: "https://your-milvus-endpoint.com/", + required: true, + type: "text", + }, + ], }; export const getVectorStoreProviderLogoAndName = (providerValue: string): { logo: string; displayName: string } => { From 0ee924f91e67381bdc5d886768c3596211684391 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 15 Dec 2025 18:34:34 -0800 Subject: [PATCH 018/121] Adding svg --- ui/litellm-dashboard/public/assets/logos/milvus.svg | 1 + 1 file changed, 1 insertion(+) create mode 100644 ui/litellm-dashboard/public/assets/logos/milvus.svg diff --git a/ui/litellm-dashboard/public/assets/logos/milvus.svg b/ui/litellm-dashboard/public/assets/logos/milvus.svg new file mode 100644 index 00000000000..76154467b4b --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/milvus.svg @@ -0,0 +1 @@ +milvus-horizontal-color \ No newline at end of file From 8e3eb331655502a00d8a1cff99da9a07ea627d0b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 15 Dec 2025 18:36:01 -0800 Subject: [PATCH 019/121] Tests --- .../VectorStoreForm.test.tsx | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.test.tsx diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.test.tsx new file mode 100644 index 00000000000..7bc948c9127 --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.test.tsx @@ -0,0 +1,28 @@ +import { describe, it, expect, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import VectorStoreForm from "./VectorStoreForm"; +import * as networking from "../networking"; +import { CredentialItem } from "../networking"; + +vi.mock("../networking"); + +describe("VectorStoreForm", () => { + it("should render the form when visible", () => { + const mockOnCancel = vi.fn(); + const mockOnSuccess = vi.fn(); + const mockAccessToken = "test-token"; + const mockCredentials: CredentialItem[] = []; + + render( + , + ); + + expect(screen.getByText("Add New Vector Store")).toBeInTheDocument(); + }); +}); From be4f11eaccea5191e6bea84f74f976abb399253d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 15 Dec 2025 18:38:41 -0800 Subject: [PATCH 020/121] Fixing build --- .../vector_store_management/VectorStoreForm.test.tsx | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.test.tsx index 7bc948c9127..97d18e33640 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.test.tsx @@ -1,8 +1,7 @@ -import { describe, it, expect, vi } from "vitest"; import { render, screen } from "@testing-library/react"; -import VectorStoreForm from "./VectorStoreForm"; -import * as networking from "../networking"; +import { describe, expect, it, vi } from "vitest"; import { CredentialItem } from "../networking"; +import VectorStoreForm from "./VectorStoreForm"; vi.mock("../networking"); From edfb4148dcf280146ed9c9359c823c35504b68a4 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 16 Dec 2025 08:13:54 +0530 Subject: [PATCH 021/121] Add workflow to create daily staging branches (#18020) Co-authored-by: Cursor Agent --- .../workflows/create_daily_staging_branch.yml | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 .github/workflows/create_daily_staging_branch.yml diff --git a/.github/workflows/create_daily_staging_branch.yml b/.github/workflows/create_daily_staging_branch.yml new file mode 100644 index 00000000000..a97cf6f9740 --- /dev/null +++ b/.github/workflows/create_daily_staging_branch.yml @@ -0,0 +1,43 @@ +name: Create Daily Staging Branch + +on: + schedule: + - cron: '0 0 * * *' # Runs daily at midnight UTC + workflow_dispatch: # Allow manual trigger + +jobs: + create-staging-branch: + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v3 + with: + fetch-depth: 0 + + - name: Create daily staging branch + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + # Configure Git user + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + + # Generate branch name with MM_DD_YYYY format + BRANCH_NAME="litellm_staging_$(date +'%m_%d_%Y')" + echo "Creating branch: $BRANCH_NAME" + + # Fetch all branches + git fetch --all + + # Check if the branch already exists + if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then + echo "Branch $BRANCH_NAME already exists. Skipping creation." + else + echo "Creating new branch: $BRANCH_NAME" + # Create the new branch from main + git checkout -b $BRANCH_NAME origin/main + # Push the new branch + git push origin $BRANCH_NAME + echo "Successfully created and pushed branch: $BRANCH_NAME" + fi From fc3f82b85a7c53e38984916e1943c10d2ef83a09 Mon Sep 17 00:00:00 2001 From: Damien Date: Mon, 15 Dec 2025 20:46:09 -0600 Subject: [PATCH 022/121] feat(gemini): support extra_headers in batch embeddings (#18004) * feat(vertex_ai): support extra_headers in batch embeddings * test(vertex_ai): add Gemini batch embeddings tests for custom api_base --- .../batch_embed_content_handler.py | 10 ++ litellm/main.py | 1 + .../vertex_ai/test_gemini_batch_embeddings.py | 145 ++++++++++++++++++ 3 files changed, 156 insertions(+) create mode 100644 tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 859bb0a6984..07f57a4a7f6 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -46,6 +46,7 @@ class GoogleBatchEmbeddings(VertexLLM): aembedding: Optional[bool] = False, timeout=300, client=None, + extra_headers: Optional[dict] = None, ) -> EmbeddingResponse: _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, @@ -90,6 +91,15 @@ class GoogleBatchEmbeddings(VertexLLM): headers = { "Content-Type": "application/json; charset=utf-8", } + if auth_header is not None: + if isinstance(auth_header, dict): + # For Gemini with custom api_base: auth_header is {"x-goog-api-key": "..."} + headers.update(auth_header) + else: + # For Vertex AI: auth_header is a Bearer token string + headers["Authorization"] = f"Bearer {auth_header}" + if extra_headers is not None: + headers.update(extra_headers) ## LOGGING logging_obj.pre_call( diff --git a/litellm/main.py b/litellm/main.py index b08ffd16e3d..216a3fe0ffd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4663,6 +4663,7 @@ def embedding( # noqa: PLR0915 api_key=gemini_api_key, api_base=api_base, client=client, + extra_headers=headers, ) elif custom_llm_provider == "vertex_ai": diff --git a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py new file mode 100644 index 00000000000..7047be4241b --- /dev/null +++ b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -0,0 +1,145 @@ +""" +Test Gemini batch embeddings with custom api_base and extra_headers. + +This test ensures that: +1. Authentication headers are properly included when using custom api_base +2. The extra_headers parameter is correctly passed through +3. Both dict-based auth_header (Gemini) and Bearer token (Vertex AI) are handled +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath("../../../..")) + +import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + + +def test_gemini_batch_embeddings_with_custom_api_base_and_auth_header(): + """ + Test that Gemini batch embeddings include auth_header when using custom api_base. + + This test verifies that when using Gemini embeddings with a custom api_base + (e.g., Cloudflare AI Gateway), the x-goog-api-key header is properly included + in the HTTP request. + """ + client = HTTPHandler() + + def mock_auth_token(*args, **kwargs): + return None, "test-project" + + with patch.object(client, "post") as mock_post, patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._ensure_access_token", + side_effect=mock_auth_token + ), patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._get_token_and_url" + ) as mock_get_token: + # Mock the _get_token_and_url to return auth_header dict and URL + mock_get_token.return_value = ( + {"x-goog-api-key": "test-gemini-api-key"}, + "https://gateway.ai.cloudflare.com/v1/test/noauth/google-ai-studio/v1beta" + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + { + "embeddings": { + "values": [0.1, 0.2, 0.3, 0.4, 0.5] + } + } + ] + } + mock_post.return_value = mock_response + + response = litellm.embedding( + model="gemini/text-embedding-004", + input=["Hello, world!"], + api_key="test-gemini-api-key", + api_base="https://gateway.ai.cloudflare.com/v1/test/noauth/google-ai-studio/v1beta", + client=client + ) + + # Verify the POST was called + mock_post.assert_called_once() + + # Get the headers that were passed to the POST request + call_args = mock_post.call_args + kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1] + headers = kwargs.get("headers", {}) + + # Verify auth_header is included + assert "x-goog-api-key" in headers, f"x-goog-api-key not in headers: {headers}" + assert headers["x-goog-api-key"] == "test-gemini-api-key" + + # Verify Content-Type is still present + assert "Content-Type" in headers + assert headers["Content-Type"] == "application/json; charset=utf-8" + + +def test_gemini_batch_embeddings_with_extra_headers(): + """ + Test that extra_headers parameter is properly included in the request. + + This test verifies that custom headers passed via extra_headers are + properly merged into the request headers. + """ + client = HTTPHandler() + + def mock_auth_token(*args, **kwargs): + return None, "test-project" + + with patch.object(client, "post") as mock_post, patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._ensure_access_token", + side_effect=mock_auth_token + ), patch( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler.GoogleBatchEmbeddings._get_token_and_url" + ) as mock_get_token: + # Mock the _get_token_and_url to return auth_header dict and URL + mock_get_token.return_value = ( + {"x-goog-api-key": "test-gemini-api-key"}, + "https://gateway.ai.cloudflare.com/v1/test/google-ai-studio/v1beta" + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + { + "embeddings": { + "values": [0.1, 0.2, 0.3] + } + } + ] + } + mock_post.return_value = mock_response + + response = litellm.embedding( + model="gemini/text-embedding-004", + input=["Test"], + api_key="test-gemini-api-key", + api_base="https://gateway.ai.cloudflare.com/v1/test/google-ai-studio/v1beta", + headers={"Authorization": "Bearer test-token", "X-Custom": "custom-value"}, + client=client + ) + + # Verify the POST was called + mock_post.assert_called_once() + + # Get the headers that were passed to the POST request + call_args = mock_post.call_args + kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1] + headers = kwargs.get("headers", {}) + + # Verify all headers are included + assert "x-goog-api-key" in headers + assert "Authorization" in headers + assert headers["Authorization"] == "Bearer test-token" + assert "X-Custom" in headers + assert headers["X-Custom"] == "custom-value" + From df9d3abf49bb76c7a2f94ef52bfd1cb257e9d6d0 Mon Sep 17 00:00:00 2001 From: Dmitrii Komarov Date: Tue, 16 Dec 2025 03:50:09 +0100 Subject: [PATCH 023/121] Propagate token usage when generating images with Gemini (#17987) --- .../vertex_gemini_transformation.py | 26 ++++++++++++++++++- ...rtex_ai_image_generation_transformation.py | 23 +++++++++++++++- 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index b9747652362..619bd006300 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -13,7 +13,7 @@ from litellm.types.llms.openai import ( AllMessageValues, OpenAIImageGenerationOptionalParams, ) -from litellm.types.utils import ImageObject, ImageResponse +from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -234,6 +234,27 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): return request_body + def _transform_image_usage(self, usage: dict) -> ImageUsage: + input_tokens_details = ImageUsageInputTokensDetails( + image_tokens=0, + text_tokens=0, + ) + tokens_details = usage.get("promptTokensDetails", []) + for details in tokens_details: + if isinstance(details, dict) and (modality := details.get("modality")): + token_count = details.get("tokenCount", 0) + if modality == "TEXT": + input_tokens_details.text_tokens += token_count + elif modality == "IMAGE": + input_tokens_details.image_tokens += token_count + + return ImageUsage( + input_tokens=usage.get("promptTokenCount", 0), + input_tokens_details=input_tokens_details, + output_tokens=usage.get("candidatesTokenCount", 0), + total_tokens=usage.get("totalTokenCount", 0), + ) + def transform_image_generation_response( self, model: str, @@ -276,6 +297,9 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): b64_json=inline_data["data"], url=None, )) + + if usage_metadata := response_data.get("usageMetadata", None): + model_response.usage = self._transform_image_usage(usage_metadata) return model_response diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 7cba03c38c8..b91438b3cac 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -141,7 +141,22 @@ class TestVertexAIGeminiImageGenerationConfig: ] } } - ] + ], + "usageMetadata": { + "promptTokenCount": 93, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 54, + }, + { + "modality": "IMAGE", + "tokenCount": 39, + } + ], + "candidatesTokenCount": 17, + "totalTokenCount": 110, + } } mock_response.headers = {} @@ -162,6 +177,12 @@ class TestVertexAIGeminiImageGenerationConfig: assert len(result.data) == 1 assert result.data[0].b64_json == "base64_encoded_image_data" assert result.data[0].url is None + assert result.usage.input_tokens == 93 + assert result.usage.input_tokens_details.text_tokens == 54 + assert result.usage.input_tokens_details.image_tokens == 39 + assert result.usage.output_tokens == 17 + assert result.usage.total_tokens == 110 + def test_transform_image_generation_response_multiple_images(self): """Test response transformation with multiple images""" From e83f08c24da1fff7ad5334c433eedf3dfe228f45 Mon Sep 17 00:00:00 2001 From: Doni Crosby Date: Mon, 15 Dec 2025 21:54:16 -0500 Subject: [PATCH 024/121] feat(venice.ai): add support for Venice.ai API via providers.json (#17962) --- litellm/llms/openai_like/providers.json | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index a6c19222619..2d801506d5f 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -14,5 +14,9 @@ "helicone": { "base_url": "https://ai-gateway.helicone.ai/", "api_key_env": "HELICONE_API_KEY" + }, + "veniceai": { + "base_url": "https://api.venice.ai/api/v1", + "api_key_env": "VENICE_AI_API_KEY" } } From 7c2478b70ef5be2e0596a600b3886421c6dbcc90 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 16 Dec 2025 08:35:39 +0530 Subject: [PATCH 025/121] docs: replace ghcr link with docker.litellm.ai --- deploy/charts/litellm-helm/README.md | 2 +- docs/my-website/docs/index.md | 2 +- docs/my-website/docs/observability/datadog.md | 2 +- docs/my-website/docs/proxy/configs.md | 4 +- docs/my-website/docs/proxy/deploy.md | 50 +++++++++---------- .../docs/proxy/docker_quick_start.md | 6 +-- .../docs/proxy/guardrails/pangea.md | 2 +- .../docs/proxy/shared_health_check.md | 2 +- .../secret_managers/custom_secret_manager.md | 2 +- .../docs/tutorials/elasticsearch_logging.md | 2 +- .../my-website/docs/tutorials/openai_codex.md | 2 +- .../release_notes/v1.55.8-stable/index.md | 2 +- .../my-website/release_notes/v1.57.3/index.md | 2 +- .../release_notes/v1.63.11-stable/index.md | 2 +- .../release_notes/v1.63.14/index.md | 2 +- .../release_notes/v1.65.4-stable/index.md | 2 +- .../release_notes/v1.66.0-stable/index.md | 2 +- .../release_notes/v1.67.4-stable/index.md | 2 +- .../release_notes/v1.68.0-stable/index.md | 2 +- .../release_notes/v1.69.0-stable/index.md | 2 +- .../release_notes/v1.70.1-stable/index.md | 2 +- .../release_notes/v1.71.1-stable/index.md | 2 +- .../release_notes/v1.72.0-stable/index.md | 2 +- .../release_notes/v1.72.2-stable/index.md | 2 +- .../release_notes/v1.72.6-stable/index.md | 2 +- .../release_notes/v1.73.0-stable/index.md | 2 +- .../release_notes/v1.73.6-stable/index.md | 2 +- .../release_notes/v1.74.0-stable/index.md | 2 +- .../release_notes/v1.74.15-stable/index.md | 2 +- .../release_notes/v1.74.3-stable/index.md | 2 +- .../my-website/release_notes/v1.74.7/index.md | 2 +- .../release_notes/v1.74.9-stable/index.md | 2 +- .../release_notes/v1.75.5-stable/index.md | 2 +- .../my-website/release_notes/v1.75.8/index.md | 2 +- .../release_notes/v1.76.1-stable/index.md | 2 +- .../release_notes/v1.76.3-stable/index.md | 2 +- .../release_notes/v1.77.2-stable/index.md | 2 +- .../release_notes/v1.77.3-stable/index.md | 2 +- .../release_notes/v1.77.5-stable/index.md | 2 +- .../release_notes/v1.77.7-stable/index.md | 2 +- .../release_notes/v1.78.0-stable/index.md | 2 +- .../release_notes/v1.78.5-stable/index.md | 2 +- .../release_notes/v1.79.0-stable/index.md | 2 +- .../release_notes/v1.79.1-stable/index.md | 2 +- .../release_notes/v1.79.3-stable/index.md | 2 +- .../release_notes/v1.80.0-stable/index.md | 2 +- .../release_notes/v1.80.10-stable/index.md | 2 +- .../release_notes/v1.80.5-stable/index.md | 2 +- .../release_notes/v1.80.8-stable/index.md | 2 +- docs/my-website/src/pages/index.md | 2 +- 50 files changed, 77 insertions(+), 77 deletions(-) diff --git a/deploy/charts/litellm-helm/README.md b/deploy/charts/litellm-helm/README.md index 6fdc423a177..2fa856843f3 100644 --- a/deploy/charts/litellm-helm/README.md +++ b/deploy/charts/litellm-helm/README.md @@ -29,7 +29,7 @@ If `db.useStackgresOperator` is used (not yet implemented): | `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A | | `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | | `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | -| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` | +| `image.repository` | LiteLLM Proxy image repository | `docker.litellm.ai/berriai/litellm` | | `image.pullPolicy` | LiteLLM Proxy image pull policy | `IfNotPresent` | | `image.tag` | Overrides the image tag whose default the latest version of LiteLLM at the time this chart was published. | `""` | | `imagePullSecrets` | Registry credentials for the LiteLLM and initContainer images. | `[]` | diff --git a/docs/my-website/docs/index.md b/docs/my-website/docs/index.md index f393b300f73..ba605e316d3 100644 --- a/docs/my-website/docs/index.md +++ b/docs/my-website/docs/index.md @@ -657,7 +657,7 @@ docker run \ -e AZURE_API_KEY=d6*********** \ -e AZURE_API_BASE=https://openai-***********/ \ -p 4000:4000 \ - ghcr.io/berriai/litellm:main-latest \ + docker.litellm.ai/berriai/litellm:main-latest \ --config /app/config.yaml --detailed_debug ``` diff --git a/docs/my-website/docs/observability/datadog.md b/docs/my-website/docs/observability/datadog.md index b2901650ea6..7cf91ced34c 100644 --- a/docs/my-website/docs/observability/datadog.md +++ b/docs/my-website/docs/observability/datadog.md @@ -181,7 +181,7 @@ docker run \ -e USE_DDTRACE=true \ -e USE_DDPROFILER=true \ -p 4000:4000 \ - ghcr.io/berriai/litellm:main-latest \ + docker.litellm.ai/berriai/litellm:main-latest \ --config /app/config.yaml --detailed_debug ``` diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index 77ab3158f74..ba4ca190aa9 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -655,7 +655,7 @@ docker run --name litellm-proxy \ -e LITELLM_CONFIG_BUCKET_OBJECT_KEY="> \ -e LITELLM_CONFIG_BUCKET_TYPE="gcs" \ -p 4000:4000 \ - ghcr.io/berriai/litellm-database:main-latest --detailed_debug + docker.litellm.ai/berriai/litellm-database:main-latest --detailed_debug ``` @@ -676,7 +676,7 @@ docker run --name litellm-proxy \ -e LITELLM_CONFIG_BUCKET_NAME= \ -e LITELLM_CONFIG_BUCKET_OBJECT_KEY="> \ -p 4000:4000 \ - ghcr.io/berriai/litellm-database:main-latest + docker.litellm.ai/berriai/litellm-database:main-latest ``` diff --git a/docs/my-website/docs/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index 0f0e5f678d3..abdb9aa3298 100644 --- a/docs/my-website/docs/proxy/deploy.md +++ b/docs/my-website/docs/proxy/deploy.md @@ -57,7 +57,7 @@ docker run \ -e AZURE_API_KEY=d6*********** \ -e AZURE_API_BASE=https://openai-***********/ \ -p 4000:4000 \ - ghcr.io/berriai/litellm:main-stable \ + docker.litellm.ai/berriai/litellm:main-stable \ --config /app/config.yaml --detailed_debug ``` @@ -87,12 +87,12 @@ See all supported CLI args [here](https://docs.litellm.ai/docs/proxy/cli): Here's how you can run the docker image and pass your config to `litellm` ```shell -docker run ghcr.io/berriai/litellm:main-stable --config your_config.yaml +docker run docker.litellm.ai/berriai/litellm:main-stable --config your_config.yaml ``` Here's how you can run the docker image and start litellm on port 8002 with `num_workers=8` ```shell -docker run ghcr.io/berriai/litellm:main-stable --port 8002 --num_workers 8 +docker run docker.litellm.ai/berriai/litellm:main-stable --port 8002 --num_workers 8 ``` @@ -100,7 +100,7 @@ docker run ghcr.io/berriai/litellm:main-stable --port 8002 --num_workers 8 ```shell # Use the provided base image -FROM ghcr.io/berriai/litellm:main-stable +FROM docker.litellm.ai/berriai/litellm:main-stable # Set the working directory to /app WORKDIR /app @@ -242,7 +242,7 @@ spec: spec: containers: - name: litellm - image: ghcr.io/berriai/litellm:main-stable # it is recommended to fix a version generally + image: docker.litellm.ai/berriai/litellm:main-stable # it is recommended to fix a version generally args: - "--config" - "/app/proxy_server_config.yaml" @@ -279,9 +279,9 @@ Use this when you want to use litellm helm chart as a dependency for other chart #### Step 1. Pull the litellm helm chart ```bash -helm pull oci://ghcr.io/berriai/litellm-helm +helm pull oci://docker.litellm.ai/berriai/litellm-helm -# Pulled: ghcr.io/berriai/litellm-helm:0.1.2 +# Pulled: docker.litellm.ai/berriai/litellm-helm:0.1.2 # Digest: sha256:7d3ded1c99c1597f9ad4dc49d84327cf1db6e0faa0eeea0c614be5526ae94e2a ``` @@ -340,7 +340,7 @@ Requirements: We maintain a [separate Dockerfile](https://github.com/BerriAI/litellm/pkgs/container/litellm-database) for reducing build time when running LiteLLM proxy with a connected Postgres Database ```shell -docker pull ghcr.io/berriai/litellm-database:main-stable +docker pull docker.litellm.ai/berriai/litellm-database:main-stable ``` ```shell @@ -351,7 +351,7 @@ docker run \ -e AZURE_API_KEY=d6*********** \ -e AZURE_API_BASE=https://openai-***********/ \ -p 4000:4000 \ - ghcr.io/berriai/litellm-database:main-stable \ + docker.litellm.ai/berriai/litellm-database:main-stable \ --config /app/config.yaml --detailed_debug ``` @@ -379,7 +379,7 @@ spec: spec: containers: - name: litellm-container - image: ghcr.io/berriai/litellm:main-stable + image: docker.litellm.ai/berriai/litellm:main-stable imagePullPolicy: Always env: - name: AZURE_API_KEY @@ -516,9 +516,9 @@ Use this when you want to use litellm helm chart as a dependency for other chart #### Step 1. Pull the litellm helm chart ```bash -helm pull oci://ghcr.io/berriai/litellm-helm +helm pull oci://docker.litellm.ai/berriai/litellm-helm -# Pulled: ghcr.io/berriai/litellm-helm:0.1.2 +# Pulled: docker.litellm.ai/berriai/litellm-helm:0.1.2 # Digest: sha256:7d3ded1c99c1597f9ad4dc49d84327cf1db6e0faa0eeea0c614be5526ae94e2a ``` @@ -575,7 +575,7 @@ router_settings: Start docker container with config ```shell -docker run ghcr.io/berriai/litellm:main-stable --config your_config.yaml +docker run docker.litellm.ai/berriai/litellm:main-stable --config your_config.yaml ``` ### Deploy with Database + Redis @@ -610,7 +610,7 @@ Start `litellm-database`docker container with config docker run --name litellm-proxy \ -e DATABASE_URL=postgresql://:@:/ \ -p 4000:4000 \ -ghcr.io/berriai/litellm-database:main-stable --config your_config.yaml +docker.litellm.ai/berriai/litellm-database:main-stable --config your_config.yaml ``` ### (Non Root) - without Internet Connection @@ -620,7 +620,7 @@ By default `prisma generate` downloads [prisma's engine binaries](https://www.pr Use this docker image to deploy litellm with pre-generated prisma binaries. ```bash -docker pull ghcr.io/berriai/litellm-non_root:main-stable +docker pull docker.litellm.ai/berriai/litellm-non_root:main-stable ``` [Published Docker Image link](https://github.com/BerriAI/litellm/pkgs/container/litellm-non_root) @@ -639,7 +639,7 @@ Use this, If you need to set ssl certificates for your on prem litellm proxy Pass `ssl_keyfile_path` (Path to the SSL keyfile) and `ssl_certfile_path` (Path to the SSL certfile) when starting litellm proxy ```shell -docker run ghcr.io/berriai/litellm:main-stable \ +docker run docker.litellm.ai/berriai/litellm:main-stable \ --ssl_keyfile_path ssl_test/keyfile.key \ --ssl_certfile_path ssl_test/certfile.crt ``` @@ -654,7 +654,7 @@ Step 1. Build your custom docker image with hypercorn ```shell # Use the provided base image -FROM ghcr.io/berriai/litellm:main-stable +FROM docker.litellm.ai/berriai/litellm:main-stable # Set the working directory to /app WORKDIR /app @@ -702,7 +702,7 @@ Usage Example: In this example, we set the keepalive timeout to 75 seconds. ```shell showLineNumbers title="docker run" -docker run ghcr.io/berriai/litellm:main-stable \ +docker run docker.litellm.ai/berriai/litellm:main-stable \ --keepalive_timeout 75 ``` @@ -711,7 +711,7 @@ In this example, we set the keepalive timeout to 75 seconds. ```shell showLineNumbers title="Environment Variable" export KEEPALIVE_TIMEOUT=75 -docker run ghcr.io/berriai/litellm:main-stable +docker run docker.litellm.ai/berriai/litellm:main-stable ``` @@ -722,7 +722,7 @@ Use this to mitigate memory growth by recycling workers after a fixed number of Usage Examples: ```shell showLineNumbers title="docker run (CLI flag)" -docker run ghcr.io/berriai/litellm:main-stable \ +docker run docker.litellm.ai/berriai/litellm:main-stable \ --max_requests_before_restart 10000 ``` @@ -730,7 +730,7 @@ Or set via environment variable: ```shell showLineNumbers title="Environment Variable" export MAX_REQUESTS_BEFORE_RESTART=10000 -docker run ghcr.io/berriai/litellm:main-stable +docker run docker.litellm.ai/berriai/litellm:main-stable ``` @@ -759,7 +759,7 @@ docker run --name litellm-proxy \ -e LITELLM_CONFIG_BUCKET_OBJECT_KEY="> \ -e LITELLM_CONFIG_BUCKET_TYPE="gcs" \ -p 4000:4000 \ - ghcr.io/berriai/litellm-database:main-stable --detailed_debug + docker.litellm.ai/berriai/litellm-database:main-stable --detailed_debug ``` @@ -780,7 +780,7 @@ docker run --name litellm-proxy \ -e LITELLM_CONFIG_BUCKET_NAME= \ -e LITELLM_CONFIG_BUCKET_OBJECT_KEY="> \ -p 4000:4000 \ - ghcr.io/berriai/litellm-database:main-stable + docker.litellm.ai/berriai/litellm-database:main-stable ``` @@ -907,7 +907,7 @@ Run the following command, replacing `` with the value you copied docker run --name litellm-proxy \ -e DATABASE_URL= \ -p 4000:4000 \ - ghcr.io/berriai/litellm-database:main-stable + docker.litellm.ai/berriai/litellm-database:main-stable ``` #### 4. Access the Application: @@ -986,7 +986,7 @@ services: context: . args: target: runtime - image: ghcr.io/berriai/litellm:main-stable + image: docker.litellm.ai/berriai/litellm:main-stable ports: - "4000:4000" # Map the container port to the host, change the host port if necessary volumes: diff --git a/docs/my-website/docs/proxy/docker_quick_start.md b/docs/my-website/docs/proxy/docker_quick_start.md index 35d9923e92c..efdc73de43e 100644 --- a/docs/my-website/docs/proxy/docker_quick_start.md +++ b/docs/my-website/docs/proxy/docker_quick_start.md @@ -20,7 +20,7 @@ End-to-End tutorial for LiteLLM Proxy to: ``` -docker pull ghcr.io/berriai/litellm:main-latest +docker pull docker.litellm.ai/berriai/litellm:main-latest ``` [**See all docker images**](https://github.com/orgs/BerriAI/packages) @@ -119,7 +119,7 @@ docker run \ -e AZURE_API_KEY=d6*********** \ -e AZURE_API_BASE=https://openai-***********/ \ -p 4000:4000 \ - ghcr.io/berriai/litellm:main-latest \ + docker.litellm.ai/berriai/litellm:main-latest \ --config /app/config.yaml --detailed_debug # RUNNING on http://0.0.0.0:4000 @@ -302,7 +302,7 @@ docker run \ -e AZURE_API_KEY=d6*********** \ -e AZURE_API_BASE=https://openai-***********/ \ -p 4000:4000 \ - ghcr.io/berriai/litellm:main-latest \ + docker.litellm.ai/berriai/litellm:main-latest \ --config /app/config.yaml --detailed_debug ``` diff --git a/docs/my-website/docs/proxy/guardrails/pangea.md b/docs/my-website/docs/proxy/guardrails/pangea.md index 180b9100d6b..3de5ddfa530 100644 --- a/docs/my-website/docs/proxy/guardrails/pangea.md +++ b/docs/my-website/docs/proxy/guardrails/pangea.md @@ -67,7 +67,7 @@ docker run --rm \ -e PANGEA_AI_GUARD_TOKEN=$PANGEA_AI_GUARD_TOKEN \ -e OPENAI_API_KEY=$OPENAI_API_KEY \ -v $(pwd)/config.yaml:/app/config.yaml \ - ghcr.io/berriai/litellm:main-latest \ + docker.litellm.ai/berriai/litellm:main-latest \ --config /app/config.yaml ``` diff --git a/docs/my-website/docs/proxy/shared_health_check.md b/docs/my-website/docs/proxy/shared_health_check.md index d4b70116309..c9c975c7911 100644 --- a/docs/my-website/docs/proxy/shared_health_check.md +++ b/docs/my-website/docs/proxy/shared_health_check.md @@ -269,7 +269,7 @@ spec: spec: containers: - name: litellm-proxy - image: ghcr.io/berriai/litellm:latest + image: docker.litellm.ai/berriai/litellm:latest env: - name: USE_SHARED_HEALTH_CHECK value: "true" diff --git a/docs/my-website/docs/secret_managers/custom_secret_manager.md b/docs/my-website/docs/secret_managers/custom_secret_manager.md index c51eeeb0727..a6a91a0336d 100644 --- a/docs/my-website/docs/secret_managers/custom_secret_manager.md +++ b/docs/my-website/docs/secret_managers/custom_secret_manager.md @@ -76,7 +76,7 @@ docker run -d \ --name litellm-proxy \ -v $(pwd)/config.yaml:/app/config.yaml \ -v $(pwd)/my_secret_manager.py:/app/my_secret_manager.py \ - ghcr.io/berriai/litellm:main-latest \ + docker.litellm.ai/berriai/litellm:main-latest \ --config /app/config.yaml \ --port 4000 \ --detailed_debug diff --git a/docs/my-website/docs/tutorials/elasticsearch_logging.md b/docs/my-website/docs/tutorials/elasticsearch_logging.md index eabd47f095d..85a9f1452d7 100644 --- a/docs/my-website/docs/tutorials/elasticsearch_logging.md +++ b/docs/my-website/docs/tutorials/elasticsearch_logging.md @@ -221,7 +221,7 @@ services: - elasticsearch litellm: - image: ghcr.io/berriai/litellm:main-latest + image: docker.litellm.ai/berriai/litellm:main-latest ports: - "4000:4000" environment: diff --git a/docs/my-website/docs/tutorials/openai_codex.md b/docs/my-website/docs/tutorials/openai_codex.md index 41416f85159..563d6559ca5 100644 --- a/docs/my-website/docs/tutorials/openai_codex.md +++ b/docs/my-website/docs/tutorials/openai_codex.md @@ -53,7 +53,7 @@ yarn global add @openai/codex docker run \ -v $(pwd)/litellm_config.yaml:/app/config.yaml \ -p 4000:4000 \ - ghcr.io/berriai/litellm:main-latest \ + docker.litellm.ai/berriai/litellm:main-latest \ --config /app/config.yaml ``` diff --git a/docs/my-website/release_notes/v1.55.8-stable/index.md b/docs/my-website/release_notes/v1.55.8-stable/index.md index 38c78eb5372..bf239e0889d 100644 --- a/docs/my-website/release_notes/v1.55.8-stable/index.md +++ b/docs/my-website/release_notes/v1.55.8-stable/index.md @@ -53,7 +53,7 @@ Send LLM usage (spend, tokens) data to [Azure Data Lake](https://learn.microsoft docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:litellm_stable_release_branch-v1.55.8-stable +docker.litellm.ai/berriai/litellm:litellm_stable_release_branch-v1.55.8-stable ``` ## Get Daily Updates diff --git a/docs/my-website/release_notes/v1.57.3/index.md b/docs/my-website/release_notes/v1.57.3/index.md index ab1154a0a8c..bbffa990b32 100644 --- a/docs/my-website/release_notes/v1.57.3/index.md +++ b/docs/my-website/release_notes/v1.57.3/index.md @@ -39,7 +39,7 @@ Instead of `apt-get` use `apk`, the base litellm image will no longer have `apt- **You are only impacted if you use `apt-get` in your Dockerfile** ```shell # Use the provided base image -FROM ghcr.io/berriai/litellm:main-latest +FROM docker.litellm.ai/berriai/litellm:main-latest # Set the working directory WORKDIR /app diff --git a/docs/my-website/release_notes/v1.63.11-stable/index.md b/docs/my-website/release_notes/v1.63.11-stable/index.md index 882747a07b3..3273f9a8e06 100644 --- a/docs/my-website/release_notes/v1.63.11-stable/index.md +++ b/docs/my-website/release_notes/v1.63.11-stable/index.md @@ -36,7 +36,7 @@ This release is primarily focused on: docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.63.11-stable +docker.litellm.ai/berriai/litellm:main-v1.63.11-stable ``` ## Demo Instance diff --git a/docs/my-website/release_notes/v1.63.14/index.md b/docs/my-website/release_notes/v1.63.14/index.md index ff2630468c5..1ac713fc2d5 100644 --- a/docs/my-website/release_notes/v1.63.14/index.md +++ b/docs/my-website/release_notes/v1.63.14/index.md @@ -32,7 +32,7 @@ This release brings: docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.63.14-stable.patch1 +docker.litellm.ai/berriai/litellm:main-v1.63.14-stable.patch1 ``` ## Demo Instance diff --git a/docs/my-website/release_notes/v1.65.4-stable/index.md b/docs/my-website/release_notes/v1.65.4-stable/index.md index 872024a47ab..80d703e1116 100644 --- a/docs/my-website/release_notes/v1.65.4-stable/index.md +++ b/docs/my-website/release_notes/v1.65.4-stable/index.md @@ -29,7 +29,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.65.4-stable +docker.litellm.ai/berriai/litellm:main-v1.65.4-stable ``` diff --git a/docs/my-website/release_notes/v1.66.0-stable/index.md b/docs/my-website/release_notes/v1.66.0-stable/index.md index 939322e0317..693cd7fc5ac 100644 --- a/docs/my-website/release_notes/v1.66.0-stable/index.md +++ b/docs/my-website/release_notes/v1.66.0-stable/index.md @@ -29,7 +29,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.66.0-stable +docker.litellm.ai/berriai/litellm:main-v1.66.0-stable ``` diff --git a/docs/my-website/release_notes/v1.67.4-stable/index.md b/docs/my-website/release_notes/v1.67.4-stable/index.md index 93a27155d2b..f61c99f7d02 100644 --- a/docs/my-website/release_notes/v1.67.4-stable/index.md +++ b/docs/my-website/release_notes/v1.67.4-stable/index.md @@ -30,7 +30,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.67.4-stable +docker.litellm.ai/berriai/litellm:main-v1.67.4-stable ``` diff --git a/docs/my-website/release_notes/v1.68.0-stable/index.md b/docs/my-website/release_notes/v1.68.0-stable/index.md index 4d456d9c853..f3e7fa27427 100644 --- a/docs/my-website/release_notes/v1.68.0-stable/index.md +++ b/docs/my-website/release_notes/v1.68.0-stable/index.md @@ -29,7 +29,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.68.0-stable +docker.litellm.ai/berriai/litellm:main-v1.68.0-stable ``` diff --git a/docs/my-website/release_notes/v1.69.0-stable/index.md b/docs/my-website/release_notes/v1.69.0-stable/index.md index 3f8ce7a29c4..f3f094e5403 100644 --- a/docs/my-website/release_notes/v1.69.0-stable/index.md +++ b/docs/my-website/release_notes/v1.69.0-stable/index.md @@ -29,7 +29,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.69.0-stable +docker.litellm.ai/berriai/litellm:main-v1.69.0-stable ``` diff --git a/docs/my-website/release_notes/v1.70.1-stable/index.md b/docs/my-website/release_notes/v1.70.1-stable/index.md index c55ac8b9c61..5d4bde0f6a0 100644 --- a/docs/my-website/release_notes/v1.70.1-stable/index.md +++ b/docs/my-website/release_notes/v1.70.1-stable/index.md @@ -30,7 +30,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.70.1-stable +docker.litellm.ai/berriai/litellm:main-v1.70.1-stable ``` diff --git a/docs/my-website/release_notes/v1.71.1-stable/index.md b/docs/my-website/release_notes/v1.71.1-stable/index.md index 2d21d49171b..bd37183455d 100644 --- a/docs/my-website/release_notes/v1.71.1-stable/index.md +++ b/docs/my-website/release_notes/v1.71.1-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.71.1-stable +docker.litellm.ai/berriai/litellm:main-v1.71.1-stable ``` diff --git a/docs/my-website/release_notes/v1.72.0-stable/index.md b/docs/my-website/release_notes/v1.72.0-stable/index.md index 47bc19e8aa8..fe235cf07b1 100644 --- a/docs/my-website/release_notes/v1.72.0-stable/index.md +++ b/docs/my-website/release_notes/v1.72.0-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.72.0-stable +docker.litellm.ai/berriai/litellm:main-v1.72.0-stable ``` diff --git a/docs/my-website/release_notes/v1.72.2-stable/index.md b/docs/my-website/release_notes/v1.72.2-stable/index.md index 023180f9758..36d01c131c7 100644 --- a/docs/my-website/release_notes/v1.72.2-stable/index.md +++ b/docs/my-website/release_notes/v1.72.2-stable/index.md @@ -29,7 +29,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.72.2-stable +docker.litellm.ai/berriai/litellm:main-v1.72.2-stable ``` diff --git a/docs/my-website/release_notes/v1.72.6-stable/index.md b/docs/my-website/release_notes/v1.72.6-stable/index.md index 5603548364f..a20488e2318 100644 --- a/docs/my-website/release_notes/v1.72.6-stable/index.md +++ b/docs/my-website/release_notes/v1.72.6-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run -e STORE_MODEL_IN_DB=True -p 4000:4000 -ghcr.io/berriai/litellm:main-v1.72.6-stable +docker.litellm.ai/berriai/litellm:main-v1.72.6-stable ``` diff --git a/docs/my-website/release_notes/v1.73.0-stable/index.md b/docs/my-website/release_notes/v1.73.0-stable/index.md index 307fecc36dd..802c5ac028b 100644 --- a/docs/my-website/release_notes/v1.73.0-stable/index.md +++ b/docs/my-website/release_notes/v1.73.0-stable/index.md @@ -37,7 +37,7 @@ The `non-root` docker image has a known issue around the UI not loading. If you docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.73.0-stable +docker.litellm.ai/berriai/litellm:v1.73.0-stable ``` diff --git a/docs/my-website/release_notes/v1.73.6-stable/index.md b/docs/my-website/release_notes/v1.73.6-stable/index.md index b03380f9b2b..da748c5c99f 100644 --- a/docs/my-website/release_notes/v1.73.6-stable/index.md +++ b/docs/my-website/release_notes/v1.73.6-stable/index.md @@ -29,7 +29,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.73.6-stable.patch.1 +docker.litellm.ai/berriai/litellm:v1.73.6-stable.patch.1 ``` diff --git a/docs/my-website/release_notes/v1.74.0-stable/index.md b/docs/my-website/release_notes/v1.74.0-stable/index.md index e49c2b4f620..ee39c0a26a8 100644 --- a/docs/my-website/release_notes/v1.74.0-stable/index.md +++ b/docs/my-website/release_notes/v1.74.0-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.74.0-stable +docker.litellm.ai/berriai/litellm:v1.74.0-stable ``` diff --git a/docs/my-website/release_notes/v1.74.15-stable/index.md b/docs/my-website/release_notes/v1.74.15-stable/index.md index 9807a00b7e7..c0facf8afb0 100644 --- a/docs/my-website/release_notes/v1.74.15-stable/index.md +++ b/docs/my-website/release_notes/v1.74.15-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.74.15-stable +docker.litellm.ai/berriai/litellm:v1.74.15-stable ``` diff --git a/docs/my-website/release_notes/v1.74.3-stable/index.md b/docs/my-website/release_notes/v1.74.3-stable/index.md index 167d81e52af..05386172e71 100644 --- a/docs/my-website/release_notes/v1.74.3-stable/index.md +++ b/docs/my-website/release_notes/v1.74.3-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.74.3-stable +docker.litellm.ai/berriai/litellm:v1.74.3-stable ``` diff --git a/docs/my-website/release_notes/v1.74.7/index.md b/docs/my-website/release_notes/v1.74.7/index.md index 7d7a568e13f..10fbd21b498 100644 --- a/docs/my-website/release_notes/v1.74.7/index.md +++ b/docs/my-website/release_notes/v1.74.7/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.74.7-stable.patch.1 +docker.litellm.ai/berriai/litellm:v1.74.7-stable.patch.1 ``` diff --git a/docs/my-website/release_notes/v1.74.9-stable/index.md b/docs/my-website/release_notes/v1.74.9-stable/index.md index 3f100745dfe..9feed6d62e6 100644 --- a/docs/my-website/release_notes/v1.74.9-stable/index.md +++ b/docs/my-website/release_notes/v1.74.9-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.74.9-stable.patch.1 +docker.litellm.ai/berriai/litellm:v1.74.9-stable.patch.1 ``` diff --git a/docs/my-website/release_notes/v1.75.5-stable/index.md b/docs/my-website/release_notes/v1.75.5-stable/index.md index 7035d285057..043f1267fc8 100644 --- a/docs/my-website/release_notes/v1.75.5-stable/index.md +++ b/docs/my-website/release_notes/v1.75.5-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.75.5-stable +docker.litellm.ai/berriai/litellm:v1.75.5-stable ``` diff --git a/docs/my-website/release_notes/v1.75.8/index.md b/docs/my-website/release_notes/v1.75.8/index.md index d7d4f37c4ee..3db1fe4b2cd 100644 --- a/docs/my-website/release_notes/v1.75.8/index.md +++ b/docs/my-website/release_notes/v1.75.8/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.75.8-stable +docker.litellm.ai/berriai/litellm:v1.75.8-stable ``` diff --git a/docs/my-website/release_notes/v1.76.1-stable/index.md b/docs/my-website/release_notes/v1.76.1-stable/index.md index 4437b7f5799..f458dfde6d4 100644 --- a/docs/my-website/release_notes/v1.76.1-stable/index.md +++ b/docs/my-website/release_notes/v1.76.1-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.76.1 +docker.litellm.ai/berriai/litellm:v1.76.1 ``` diff --git a/docs/my-website/release_notes/v1.76.3-stable/index.md b/docs/my-website/release_notes/v1.76.3-stable/index.md index 6b40e4f5b35..9763a57975b 100644 --- a/docs/my-website/release_notes/v1.76.3-stable/index.md +++ b/docs/my-website/release_notes/v1.76.3-stable/index.md @@ -35,7 +35,7 @@ This release has a known issue where startup is leading to Out of Memory errors docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.76.3 +docker.litellm.ai/berriai/litellm:v1.76.3 ``` diff --git a/docs/my-website/release_notes/v1.77.2-stable/index.md b/docs/my-website/release_notes/v1.77.2-stable/index.md index fdd80693d05..4f732a1604d 100644 --- a/docs/my-website/release_notes/v1.77.2-stable/index.md +++ b/docs/my-website/release_notes/v1.77.2-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:main-v1.77.2-stable +docker.litellm.ai/berriai/litellm:main-v1.77.2-stable ``` diff --git a/docs/my-website/release_notes/v1.77.3-stable/index.md b/docs/my-website/release_notes/v1.77.3-stable/index.md index c7c17e5baee..11b82c4c834 100644 --- a/docs/my-website/release_notes/v1.77.3-stable/index.md +++ b/docs/my-website/release_notes/v1.77.3-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.77.3-stable +docker.litellm.ai/berriai/litellm:v1.77.3-stable ``` diff --git a/docs/my-website/release_notes/v1.77.5-stable/index.md b/docs/my-website/release_notes/v1.77.5-stable/index.md index 6843800ee6d..8e59ea92cc2 100644 --- a/docs/my-website/release_notes/v1.77.5-stable/index.md +++ b/docs/my-website/release_notes/v1.77.5-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.77.5-stable +docker.litellm.ai/berriai/litellm:v1.77.5-stable ``` diff --git a/docs/my-website/release_notes/v1.77.7-stable/index.md b/docs/my-website/release_notes/v1.77.7-stable/index.md index 62d9a2eee4f..b4df447f334 100644 --- a/docs/my-website/release_notes/v1.77.7-stable/index.md +++ b/docs/my-website/release_notes/v1.77.7-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.77.7.rc.1 +docker.litellm.ai/berriai/litellm:v1.77.7.rc.1 ``` diff --git a/docs/my-website/release_notes/v1.78.0-stable/index.md b/docs/my-website/release_notes/v1.78.0-stable/index.md index 7f6c5ba1e08..8322f0479c5 100644 --- a/docs/my-website/release_notes/v1.78.0-stable/index.md +++ b/docs/my-website/release_notes/v1.78.0-stable/index.md @@ -28,7 +28,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.78.0-stable +docker.litellm.ai/berriai/litellm:v1.78.0-stable ``` diff --git a/docs/my-website/release_notes/v1.78.5-stable/index.md b/docs/my-website/release_notes/v1.78.5-stable/index.md index af1fd359fa2..2bcdfab472c 100644 --- a/docs/my-website/release_notes/v1.78.5-stable/index.md +++ b/docs/my-website/release_notes/v1.78.5-stable/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.78.5-stable +docker.litellm.ai/berriai/litellm:v1.78.5-stable ``` diff --git a/docs/my-website/release_notes/v1.79.0-stable/index.md b/docs/my-website/release_notes/v1.79.0-stable/index.md index 8327f4b6178..4bb7094a3fc 100644 --- a/docs/my-website/release_notes/v1.79.0-stable/index.md +++ b/docs/my-website/release_notes/v1.79.0-stable/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.79.0-stable +docker.litellm.ai/berriai/litellm:v1.79.0-stable ``` diff --git a/docs/my-website/release_notes/v1.79.1-stable/index.md b/docs/my-website/release_notes/v1.79.1-stable/index.md index ea8cfeae740..19fc7f9f3ff 100644 --- a/docs/my-website/release_notes/v1.79.1-stable/index.md +++ b/docs/my-website/release_notes/v1.79.1-stable/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.79.1-stable +docker.litellm.ai/berriai/litellm:v1.79.1-stable ``` diff --git a/docs/my-website/release_notes/v1.79.3-stable/index.md b/docs/my-website/release_notes/v1.79.3-stable/index.md index c4f3ba1e017..542f88787e0 100644 --- a/docs/my-website/release_notes/v1.79.3-stable/index.md +++ b/docs/my-website/release_notes/v1.79.3-stable/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.79.3-stable +docker.litellm.ai/berriai/litellm:v1.79.3-stable ``` diff --git a/docs/my-website/release_notes/v1.80.0-stable/index.md b/docs/my-website/release_notes/v1.80.0-stable/index.md index 17fcf6646ed..d0cf28a5c58 100644 --- a/docs/my-website/release_notes/v1.80.0-stable/index.md +++ b/docs/my-website/release_notes/v1.80.0-stable/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.80.0-stable +docker.litellm.ai/berriai/litellm:v1.80.0-stable ``` diff --git a/docs/my-website/release_notes/v1.80.10-stable/index.md b/docs/my-website/release_notes/v1.80.10-stable/index.md index 1b0a9866fae..2290c06de53 100644 --- a/docs/my-website/release_notes/v1.80.10-stable/index.md +++ b/docs/my-website/release_notes/v1.80.10-stable/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.80.10.rc.1 +docker.litellm.ai/berriai/litellm:v1.80.10.rc.1 ``` diff --git a/docs/my-website/release_notes/v1.80.5-stable/index.md b/docs/my-website/release_notes/v1.80.5-stable/index.md index 598fa47f223..9c769f8996f 100644 --- a/docs/my-website/release_notes/v1.80.5-stable/index.md +++ b/docs/my-website/release_notes/v1.80.5-stable/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.80.5-stable +docker.litellm.ai/berriai/litellm:v1.80.5-stable ``` diff --git a/docs/my-website/release_notes/v1.80.8-stable/index.md b/docs/my-website/release_notes/v1.80.8-stable/index.md index 29075a9594f..106c594968f 100644 --- a/docs/my-website/release_notes/v1.80.8-stable/index.md +++ b/docs/my-website/release_notes/v1.80.8-stable/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.80.8-stable +docker.litellm.ai/berriai/litellm:v1.80.8-stable ``` diff --git a/docs/my-website/src/pages/index.md b/docs/my-website/src/pages/index.md index 1dc2995c5fe..91215b33c5d 100644 --- a/docs/my-website/src/pages/index.md +++ b/docs/my-website/src/pages/index.md @@ -604,7 +604,7 @@ docker run \ -e AZURE_API_KEY=d6*********** \ -e AZURE_API_BASE=https://openai-***********/ \ -p 4000:4000 \ - ghcr.io/berriai/litellm:main-latest \ + docker.litellm.ai/berriai/litellm:main-latest \ --config /app/config.yaml --detailed_debug ``` From a0754f1c88acf5f64cc69e46b103340f08c05c60 Mon Sep 17 00:00:00 2001 From: kothamah <104782493+kothamah@users.noreply.github.com> Date: Mon, 15 Dec 2025 22:11:13 -0500 Subject: [PATCH 026/121] Litellm bedrock guardrails block precedence over masking (#17968) * prioritized bedrock guardrail blocking by removing early return based on masking flags When mask_request_content: true or mask_response_content: true, the method immediately returning False. The Result: Even when Bedrock Guardrails returned action: "BLOCKED" for dangerous content, LiteLLM would not raise an exception and allowing the content through the response. So removed that early condition which will return true for the blocked actions based on guardrails. * Added test case for bedrock guardrail block content precedence --- .../guardrail_hooks/bedrock_guardrails.py | 10 +-- .../test_bedrock_guardrails.py | 88 +++++++++++++++++++ 2 files changed, 93 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index fb14ccce50c..62c997659bd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -605,13 +605,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions. - If `self.mask_request_content` or `self.mask_response_content` is set to `True`, then use the output from the guardrail to mask the request or response content. + If `self.mask_request_content` or `self.mask_response_content` is set to `True`, + then use the output from the guardrail to mask the request or response content. + + However, even with masking enabled, content with action="BLOCKED" should still + raise an exception, only content with action="ANONYMIZED" should be masked. """ - # if user opted into masking, return False. since we'll use the masked output from the guardrail - if self.mask_request_content or self.mask_response_content: - return False - # if no intervention, return False if response.get("action") != "GUARDRAIL_INTERVENED": return False diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 69b0bb27b4b..84d320a0a27 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1101,3 +1101,91 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): # Verify that the Bedrock API was NOT called since there's no text to process mock_api_request.assert_not_called() print("✅ apply_guardrail with tool_calls test passed - no API call made") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): + """Test that BLOCKED content raises exception even when masking is enabled + + This test verifies the bug fix where previously mask_request_content=True or + mask_response_content=True would bypass all BLOCKED content checks. Now it + properly distinguishes between BLOCKED (raise exception) and ANONYMIZED (apply masking). + """ + + # Create guardrail with masking enabled + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + mask_request_content=True, # Masking enabled + mask_response_content=True, # Masking enabled + ) + + # Mock Bedrock response with BLOCKED content (hate speech) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", # Should raise exception + } + ] + }, + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "NAME", + "match": "John Doe", + "action": "ANONYMIZED", # Should be masked + } + ] + }, + } + ], + "outputs": [{"text": "Content blocked due to policy violation"}], + } + + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = blocked_response + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Test message with PII and hate speech"}, + ], + } + + # Mock AWS-related methods + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + mock_post.return_value = mock_bedrock_response + + # Should raise HTTPException for BLOCKED content + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, + ) + + # Verify exception details + assert exc_info.value.status_code == 400 + assert "Violated guardrail policy" in str(exc_info.value.detail) + + print("✅ BLOCKED content with masking enabled raises exception correctly") + From 0abe5cdce98e4540d402c5607429b4d8182a1f83 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 16 Dec 2025 08:41:36 +0530 Subject: [PATCH 027/121] Revert "Litellm bedrock guardrails block precedence over masking (#17968)" (#18022) This reverts commit a0754f1c88acf5f64cc69e46b103340f08c05c60. --- .../guardrail_hooks/bedrock_guardrails.py | 10 +-- .../test_bedrock_guardrails.py | 88 ------------------- 2 files changed, 5 insertions(+), 93 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 62c997659bd..fb14ccce50c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -605,13 +605,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions. - If `self.mask_request_content` or `self.mask_response_content` is set to `True`, - then use the output from the guardrail to mask the request or response content. - - However, even with masking enabled, content with action="BLOCKED" should still - raise an exception, only content with action="ANONYMIZED" should be masked. + If `self.mask_request_content` or `self.mask_response_content` is set to `True`, then use the output from the guardrail to mask the request or response content. """ + # if user opted into masking, return False. since we'll use the masked output from the guardrail + if self.mask_request_content or self.mask_response_content: + return False + # if no intervention, return False if response.get("action") != "GUARDRAIL_INTERVENED": return False diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 84d320a0a27..69b0bb27b4b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1101,91 +1101,3 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): # Verify that the Bedrock API was NOT called since there's no text to process mock_api_request.assert_not_called() print("✅ apply_guardrail with tool_calls test passed - no API call made") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): - """Test that BLOCKED content raises exception even when masking is enabled - - This test verifies the bug fix where previously mask_request_content=True or - mask_response_content=True would bypass all BLOCKED content checks. Now it - properly distinguishes between BLOCKED (raise exception) and ANONYMIZED (apply masking). - """ - - # Create guardrail with masking enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - mask_request_content=True, # Masking enabled - mask_response_content=True, # Masking enabled - ) - - # Mock Bedrock response with BLOCKED content (hate speech) - blocked_response = { - "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "contentPolicy": { - "filters": [ - { - "type": "HATE", - "confidence": "HIGH", - "action": "BLOCKED", # Should raise exception - } - ] - }, - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "NAME", - "match": "John Doe", - "action": "ANONYMIZED", # Should be masked - } - ] - }, - } - ], - "outputs": [{"text": "Content blocked due to policy violation"}], - } - - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = blocked_response - - # Mock credentials - mock_credentials = MagicMock() - mock_credentials.access_key = "test-access-key" - mock_credentials.secret_key = "test-secret-key" - mock_credentials.token = None - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Test message with PII and hate speech"}, - ], - } - - # Mock AWS-related methods - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ): - mock_post.return_value = mock_bedrock_response - - # Should raise HTTPException for BLOCKED content - with pytest.raises(HTTPException) as exc_info: - await guardrail.make_bedrock_api_request( - source="INPUT", - messages=request_data.get("messages"), - request_data=request_data, - ) - - # Verify exception details - assert exc_info.value.status_code == 400 - assert "Violated guardrail policy" in str(exc_info.value.detail) - - print("✅ BLOCKED content with masking enabled raises exception correctly") - From f58b76aee8f97fe223d6ba40d47159a8d04185ed Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 16 Dec 2025 08:42:10 +0530 Subject: [PATCH 028/121] =?UTF-8?q?Revert=20"Revert=20"Litellm=20bedrock?= =?UTF-8?q?=20guardrails=20block=20precedence=20over=20masking=20(#17?= =?UTF-8?q?=E2=80=A6"=20(#18023)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 0abe5cdce98e4540d402c5607429b4d8182a1f83. --- .../guardrail_hooks/bedrock_guardrails.py | 10 +-- .../test_bedrock_guardrails.py | 88 +++++++++++++++++++ 2 files changed, 93 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index fb14ccce50c..62c997659bd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -605,13 +605,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions. - If `self.mask_request_content` or `self.mask_response_content` is set to `True`, then use the output from the guardrail to mask the request or response content. + If `self.mask_request_content` or `self.mask_response_content` is set to `True`, + then use the output from the guardrail to mask the request or response content. + + However, even with masking enabled, content with action="BLOCKED" should still + raise an exception, only content with action="ANONYMIZED" should be masked. """ - # if user opted into masking, return False. since we'll use the masked output from the guardrail - if self.mask_request_content or self.mask_response_content: - return False - # if no intervention, return False if response.get("action") != "GUARDRAIL_INTERVENED": return False diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 69b0bb27b4b..84d320a0a27 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1101,3 +1101,91 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): # Verify that the Bedrock API was NOT called since there's no text to process mock_api_request.assert_not_called() print("✅ apply_guardrail with tool_calls test passed - no API call made") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): + """Test that BLOCKED content raises exception even when masking is enabled + + This test verifies the bug fix where previously mask_request_content=True or + mask_response_content=True would bypass all BLOCKED content checks. Now it + properly distinguishes between BLOCKED (raise exception) and ANONYMIZED (apply masking). + """ + + # Create guardrail with masking enabled + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + mask_request_content=True, # Masking enabled + mask_response_content=True, # Masking enabled + ) + + # Mock Bedrock response with BLOCKED content (hate speech) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", # Should raise exception + } + ] + }, + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "NAME", + "match": "John Doe", + "action": "ANONYMIZED", # Should be masked + } + ] + }, + } + ], + "outputs": [{"text": "Content blocked due to policy violation"}], + } + + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = blocked_response + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Test message with PII and hate speech"}, + ], + } + + # Mock AWS-related methods + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + mock_post.return_value = mock_bedrock_response + + # Should raise HTTPException for BLOCKED content + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, + ) + + # Verify exception details + assert exc_info.value.status_code == 400 + assert "Violated guardrail policy" in str(exc_info.value.detail) + + print("✅ BLOCKED content with masking enabled raises exception correctly") + From ae7c994526eed0cc2b13b1cc21e33ce37799bfcb Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Tue, 16 Dec 2025 11:12:57 +0800 Subject: [PATCH 029/121] fix(proxy): extract model from vertex ai passthrough url pattern (#17970) extract model id from vertex ai passthrough routes that follow the pattern: /vertex_ai/*/models/{model_id}:* the model extraction now handles vertex ai routes by regex matching the model segment from the url path, which allows proper model identification for authentication and authorization in proxy pass-through endpoints. adds comprehensive test coverage for vertex ai model extraction including: - various vertex api versions (v1, v1beta1) - different locations (us-central1, asia-southeast1) - model names with special suffixes (gemini-1.5-pro, gemini-2.0-flash) - precedence verification (request body model over url) - non-vertex route isolation --- litellm/proxy/auth/auth_utils.py | 8 + tests/local_testing/test_auth_utils.py | 53 +++ .../test_vertex_passthrough_auth.py | 385 ++++++++++++++++++ 3 files changed, 446 insertions(+) create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_auth.py diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index c4d0d2f8f1c..7a71af1da5c 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -616,6 +616,14 @@ def get_model_from_request( if match: model = match.group(1) + # If still not found, extract from Vertex AI passthrough route + # Pattern: /vertex_ai/.../models/{model_id}:* + # Example: /vertex_ai/v1/.../models/gemini-1.5-pro:generateContent + if model is None and "/vertex" in route.lower(): + vertex_match = re.search(r"/models/([^/:]+)", route) + if vertex_match: + model = vertex_match.group(1) + return model diff --git a/tests/local_testing/test_auth_utils.py b/tests/local_testing/test_auth_utils.py index 11261592c32..72f799a6cf0 100644 --- a/tests/local_testing/test_auth_utils.py +++ b/tests/local_testing/test_auth_utils.py @@ -311,3 +311,56 @@ def test_get_internal_user_header_from_mapping_no_internal_returns_none(): single_mapping = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(single_mapping) assert result is None + + +@pytest.mark.parametrize( + "request_data, route, expected_model", + [ + # Vertex AI passthrough URL patterns + ( + {}, + "/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + "gemini-1.5-pro" + ), + ( + {}, + "/vertex_ai/v1beta1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.0-pro:streamGenerateContent", + "gemini-1.0-pro" + ), + ( + {}, + "/vertex_ai/v1/projects/my-project/locations/asia-southeast1/publishers/google/models/gemini-2.0-flash:generateContent", + "gemini-2.0-flash" + ), + # Model without method suffix (no colon) - should still extract + ( + {}, + "/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-pro", + "gemini-pro" # Should match even without colon + ), + # Request body model takes precedence over URL + ( + {"model": "gpt-4o"}, + "/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + "gpt-4o" + ), + # Non-vertex route should not extract from vertex pattern + ( + {}, + "/openai/v1/chat/completions", + None + ), + # Azure deployment pattern should still work + ( + {}, + "/openai/deployments/my-deployment/chat/completions", + "my-deployment" + ), + ], +) +def test_get_model_from_request_vertex_ai_passthrough(request_data, route, expected_model): + """Test that get_model_from_request correctly extracts Vertex AI model from URL""" + from litellm.proxy.auth.auth_utils import get_model_from_request + + model = get_model_from_request(request_data, route) + assert model == expected_model diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_auth.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_auth.py new file mode 100644 index 00000000000..71f2883e8cd --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_auth.py @@ -0,0 +1,385 @@ + +import pytest +from unittest.mock import MagicMock, AsyncMock, patch +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _base_vertex_proxy_route +from litellm.proxy._types import ProxyException, ProxyErrorTypes +from fastapi import status + + +@pytest.mark.asyncio +async def test_vertex_passthrough_model_access_allowed_exact_match(): + """Verify that users can access Vertex models they have permission for (exact match)""" + # Setup mocks + mock_request = MagicMock() + mock_response = MagicMock() + mock_handler = MagicMock() + + # Mock user API key with access to gemini-1.5-pro + mock_user_api_key = MagicMock() + mock_user_api_key.models = ["gemini-1.5-pro"] + mock_user_api_key.team_model_aliases = None + mock_user_api_key.token = "sk-1234567890abcdef" + + # Mock router + mock_router = MagicMock() + mock_deployment = { + "litellm_params": { + "model": "vertex_ai/gemini-1.5-pro", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "use_in_pass_through": True + } + } + mock_router.get_available_deployment_for_pass_through.return_value = mock_deployment + + with patch("litellm.llms.vertex_ai.common_utils.get_vertex_model_id_from_url", return_value="gemini-1.5-pro"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_project_id_from_url", return_value="my-project"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_location_from_url", return_value="us-central1"), \ + patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.proxy_server.llm_model_list", []), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router") as mock_pt_router, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock) as mock_prep_headers, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route") as mock_create_route, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._override_vertex_params_from_router_credentials") as mock_override: + + # Setup + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ({}, "https://test.url", False, "my-project", "us-central1") + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = mock_user_api_key + mock_override.return_value = ("my-project", "us-central1") + + # Execute + await _base_vertex_proxy_route( + endpoint="v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + request=mock_request, + fastapi_response=mock_response, + get_vertex_pass_through_handler=mock_handler + ) + + # Verify - user_api_key_auth was called (which includes model permission check) + mock_auth.assert_called_once() + + +@pytest.mark.asyncio +async def test_vertex_passthrough_model_access_denied(): + """Verify that users cannot access models they don't have permission for""" + # Setup mocks + mock_request = MagicMock() + mock_response = MagicMock() + mock_handler = MagicMock() + + # Mock user API key with access only to gemini-pro, not gemini-1.5-pro + mock_user_api_key = MagicMock() + mock_user_api_key.models = ["gemini-pro"] + mock_user_api_key.team_model_aliases = None + mock_user_api_key.token = "sk-1234567890abcdef" + + with patch("litellm.llms.vertex_ai.common_utils.get_vertex_model_id_from_url", return_value="gemini-1.5-pro"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_project_id_from_url", return_value="my-project"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_location_from_url", return_value="us-central1"), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._override_vertex_params_from_router_credentials") as mock_override: + + # Setup + # Simulate access denied by raising ProxyException during user_api_key_auth + mock_auth.side_effect = ProxyException( + message="Key not allowed to access model", + type=ProxyErrorTypes.auth_error, + param="model", + code=status.HTTP_401_UNAUTHORIZED, + ) + mock_override.return_value = ("my-project", "us-central1") + + # Execute and expect exception + with pytest.raises(ProxyException) as exc_info: + await _base_vertex_proxy_route( + endpoint="v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + request=mock_request, + fastapi_response=mock_response, + get_vertex_pass_through_handler=mock_handler + ) + + # Verify + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "401" or exc_info.value.code == status.HTTP_401_UNAUTHORIZED + + +@pytest.mark.asyncio +async def test_vertex_passthrough_wildcard_access(): + """Verify that wildcard permissions (vertex_ai/*) work correctly""" + # Setup mocks + mock_request = MagicMock() + mock_response = MagicMock() + mock_handler = MagicMock() + + # Mock user API key with wildcard access to all vertex_ai models + mock_user_api_key = MagicMock() + mock_user_api_key.models = ["vertex_ai/*"] + mock_user_api_key.team_model_aliases = None + mock_user_api_key.token = "sk-1234567890abcdef" + + # Mock router + mock_router = MagicMock() + mock_deployment = { + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "use_in_pass_through": True + } + } + mock_router.get_available_deployment_for_pass_through.return_value = mock_deployment + + with patch("litellm.llms.vertex_ai.common_utils.get_vertex_model_id_from_url", return_value="gemini-2.0-flash"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_project_id_from_url", return_value="my-project"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_location_from_url", return_value="us-central1"), \ + patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.proxy_server.llm_model_list", []), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router") as mock_pt_router, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock) as mock_prep_headers, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route") as mock_create_route, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._override_vertex_params_from_router_credentials") as mock_override: + + # Setup + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ({}, "https://test.url", False, "my-project", "us-central1") + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = mock_user_api_key + mock_override.return_value = ("my-project", "us-central1") + + # Execute + await _base_vertex_proxy_route( + endpoint="v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent", + request=mock_request, + fastapi_response=mock_response, + get_vertex_pass_through_handler=mock_handler + ) + + # Verify - user_api_key_auth was called (which includes model permission check) + mock_auth.assert_called_once() + + +@pytest.mark.asyncio +async def test_vertex_passthrough_access_group(): + """Verify that access group permissions work correctly""" + # Setup mocks + mock_request = MagicMock() + mock_response = MagicMock() + mock_handler = MagicMock() + + # Mock user API key with access to a group + mock_user_api_key = MagicMock() + mock_user_api_key.models = ["production-models"] + mock_user_api_key.team_model_aliases = None + mock_user_api_key.token = "sk-1234567890abcdef" + + # Mock router with access groups + mock_router = MagicMock() + mock_deployment = { + "litellm_params": { + "model": "vertex_ai/gemini-1.5-pro", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "use_in_pass_through": True + } + } + mock_router.get_available_deployment_for_pass_through.return_value = mock_deployment + # Simulate access group matching + mock_router.get_model_access_groups.return_value = {"gemini-1.5-pro": ["production-models"]} + + with patch("litellm.llms.vertex_ai.common_utils.get_vertex_model_id_from_url", return_value="gemini-1.5-pro"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_project_id_from_url", return_value="my-project"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_location_from_url", return_value="us-central1"), \ + patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.proxy_server.llm_model_list", []), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router") as mock_pt_router, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock) as mock_prep_headers, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route") as mock_create_route, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._override_vertex_params_from_router_credentials") as mock_override: + + # Setup + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ({}, "https://test.url", False, "my-project", "us-central1") + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = mock_user_api_key + mock_override.return_value = ("my-project", "us-central1") + + # Execute + await _base_vertex_proxy_route( + endpoint="v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + request=mock_request, + fastapi_response=mock_response, + get_vertex_pass_through_handler=mock_handler + ) + + # Verify - user_api_key_auth was called (which includes model permission check) + mock_auth.assert_called_once() + + +@pytest.mark.asyncio +async def test_vertex_passthrough_team_alias(): + """Verify that team model aliases work correctly""" + # Setup mocks + mock_request = MagicMock() + mock_response = MagicMock() + mock_handler = MagicMock() + + # Mock user API key with team alias + mock_user_api_key = MagicMock() + mock_user_api_key.models = ["my-gemini"] + mock_user_api_key.team_model_aliases = {"my-gemini": "gemini-1.5-pro"} + mock_user_api_key.token = "sk-1234567890abcdef" + + # Mock router + mock_router = MagicMock() + mock_deployment = { + "litellm_params": { + "model": "vertex_ai/gemini-1.5-pro", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "use_in_pass_through": True + } + } + mock_router.get_available_deployment_for_pass_through.return_value = mock_deployment + + with patch("litellm.llms.vertex_ai.common_utils.get_vertex_model_id_from_url", return_value="gemini-1.5-pro"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_project_id_from_url", return_value="my-project"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_location_from_url", return_value="us-central1"), \ + patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.proxy_server.llm_model_list", []), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router") as mock_pt_router, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock) as mock_prep_headers, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route") as mock_create_route, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._override_vertex_params_from_router_credentials") as mock_override: + + # Setup + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ({}, "https://test.url", False, "my-project", "us-central1") + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = mock_user_api_key + mock_override.return_value = ("my-project", "us-central1") + + # Execute + await _base_vertex_proxy_route( + endpoint="v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + request=mock_request, + fastapi_response=mock_response, + get_vertex_pass_through_handler=mock_handler + ) + + # Verify - user_api_key_auth was called (which includes model permission check) + mock_auth.assert_called_once() + + +@pytest.mark.asyncio +async def test_vertex_passthrough_no_model_id(): + """Verify graceful handling when model_id cannot be extracted""" + # Setup mocks + mock_request = MagicMock() + mock_response = MagicMock() + mock_handler = MagicMock() + + # Mock user API key + mock_user_api_key = MagicMock() + mock_user_api_key.models = ["*"] + + # Mock router + mock_router = MagicMock() + + with patch("litellm.llms.vertex_ai.common_utils.get_vertex_model_id_from_url", return_value=None), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_project_id_from_url", return_value="my-project"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_location_from_url", return_value="us-central1"), \ + patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.proxy_server.llm_model_list", []), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router") as mock_pt_router, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock) as mock_prep_headers, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route") as mock_create_route, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._override_vertex_params_from_router_credentials") as mock_override: + + # Setup + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ({}, "https://test.url", False, "my-project", "us-central1") + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = mock_user_api_key + mock_override.return_value = ("my-project", "us-central1") + + # Execute - should not raise exception even though model_id is None + await _base_vertex_proxy_route( + endpoint="v1/projects/my-project/locations/us-central1/publishers/google/models/", + request=mock_request, + fastapi_response=mock_response, + get_vertex_pass_through_handler=mock_handler + ) + + # Verify - function executed successfully without errors + mock_auth.assert_called_once() + + +@pytest.mark.asyncio +async def test_vertex_passthrough_with_router_deployment(): + """Verify that permission checks don't affect existing router deployment lookup logic""" + # Setup mocks + mock_request = MagicMock() + mock_response = MagicMock() + mock_handler = MagicMock() + + # Mock user API key with wildcard access + mock_user_api_key = MagicMock() + mock_user_api_key.models = ["*"] + mock_user_api_key.team_model_aliases = None + + # Mock router with deployment + mock_router = MagicMock() + mock_deployment = { + "litellm_params": { + "model": "vertex_ai/gemini-1.5-pro", + "vertex_project": "deployment-project", + "vertex_location": "deployment-location", + "use_in_pass_through": True + } + } + mock_router.get_available_deployment_for_pass_through.return_value = mock_deployment + + with patch("litellm.llms.vertex_ai.common_utils.get_vertex_model_id_from_url", return_value="gemini-1.5-pro"), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_project_id_from_url", return_value=None), \ + patch("litellm.llms.vertex_ai.common_utils.get_vertex_location_from_url", return_value=None), \ + patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.proxy_server.llm_model_list", []), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router") as mock_pt_router, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock) as mock_prep_headers, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route") as mock_create_route, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._override_vertex_params_from_router_credentials") as mock_override: + + # Setup + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ({}, "https://test.url", False, "deployment-project", "deployment-location") + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = mock_user_api_key + mock_override.return_value = (None, None) + + # Execute + await _base_vertex_proxy_route( + endpoint="v1/projects/my-project/locations/my-location/publishers/google/models/gemini-1.5-pro:generateContent", + request=mock_request, + fastapi_response=mock_response, + get_vertex_pass_through_handler=mock_handler + ) + + # Verify + # 1. user_api_key_auth was called (which includes model permission check) + mock_auth.assert_called_once() + # 2. get_available_deployment_for_pass_through was called to get project/location + mock_router.get_available_deployment_for_pass_through.assert_called_once_with(model="gemini-1.5-pro") + # 3. Verify deployment project/location were used in auth headers + call_args = mock_prep_headers.call_args + assert call_args[1]['vertex_project'] == "deployment-project" + assert call_args[1]['vertex_location'] == "deployment-location" From 32c07113cf8a49c238144144cf07f0e2f8005e24 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 15 Dec 2025 19:14:49 -0800 Subject: [PATCH 030/121] [Feat] New Provider - VertexAI Agent Engine (#18014) * init A2AProviderConfigManager * move file * move file * add pydnatic ai folder * init providers * test_pydantic_ai_non_streaming * fix import * INIT pydantic * use_a2a_form_fields * test_vertex_agent_engine_streaming * add agent_engine * init transform for agent engine * init agent engine * VertexAgentEngineSSEStreamIterator * sample * ui add new fields * fix vertex_credentials * working SSE iterator * TestVertexAgentEngineTransformRequest * fix code QA check * Potential fix for code scanning alert no. 3923: Clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- .../llms/vertex_ai/agent_engine/__init__.py | 13 + .../vertex_ai/agent_engine/sse_iterator.py | 90 ++++ .../vertex_ai/agent_engine/transformation.py | 508 ++++++++++++++++++ litellm/llms/vertex_ai/common_utils.py | 7 +- litellm/main.py | 31 ++ ...odel_prices_and_context_window_backup.json | 2 +- .../public_endpoints/agent_create_fields.json | 23 + .../provider_create_fields.json | 4 +- tests/agent_tests/local_vertex_agent.py | 151 ++++++ .../agent_tests/test_a2a_completion_bridge.py | 76 +++ .../agent_engine/test_transformation.py | 128 +++++ 11 files changed, 1029 insertions(+), 4 deletions(-) create mode 100644 litellm/llms/vertex_ai/agent_engine/__init__.py create mode 100644 litellm/llms/vertex_ai/agent_engine/sse_iterator.py create mode 100644 litellm/llms/vertex_ai/agent_engine/transformation.py create mode 100644 tests/agent_tests/local_vertex_agent.py create mode 100644 tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py diff --git a/litellm/llms/vertex_ai/agent_engine/__init__.py b/litellm/llms/vertex_ai/agent_engine/__init__.py new file mode 100644 index 00000000000..de891f85602 --- /dev/null +++ b/litellm/llms/vertex_ai/agent_engine/__init__.py @@ -0,0 +1,13 @@ +""" +Vertex AI Agent Engine (Reasoning Engines) Provider + +Supports Vertex AI Reasoning Engines via the :query and :streamQuery endpoints. +""" + +from litellm.llms.vertex_ai.agent_engine.transformation import ( + VertexAgentEngineConfig, + VertexAgentEngineError, +) + +__all__ = ["VertexAgentEngineConfig", "VertexAgentEngineError"] + diff --git a/litellm/llms/vertex_ai/agent_engine/sse_iterator.py b/litellm/llms/vertex_ai/agent_engine/sse_iterator.py new file mode 100644 index 00000000000..06fb55e1848 --- /dev/null +++ b/litellm/llms/vertex_ai/agent_engine/sse_iterator.py @@ -0,0 +1,90 @@ +""" +SSE Stream Iterator for Vertex AI Agent Engine. + +Handles Server-Sent Events (SSE) streaming responses from Vertex AI Reasoning Engines. +""" + +from typing import Any, Union + +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.types.llms.openai import ChatCompletionUsageBlock +from litellm.types.utils import ( + Delta, + GenericStreamingChunk, + ModelResponseStream, + StreamingChoices, +) + + +class VertexAgentEngineResponseIterator(BaseModelResponseIterator): + """ + Iterator for Vertex Agent Engine SSE streaming responses. + + Uses BaseModelResponseIterator which handles sync/async iteration. + We just need to implement chunk_parser to parse Vertex Agent Engine response format. + """ + + def __init__(self, streaming_response: Any, sync_stream: bool) -> None: + super().__init__(streaming_response=streaming_response, sync_stream=sync_stream) + + def chunk_parser( + self, chunk: dict + ) -> Union[GenericStreamingChunk, ModelResponseStream]: + """ + Parse a Vertex Agent Engine response chunk into ModelResponseStream. + + Vertex Agent Engine response format: + { + "content": { + "parts": [{"text": "..."}], + "role": "model" + }, + "finish_reason": "STOP", + "usage_metadata": { + "prompt_token_count": 100, + "candidates_token_count": 50, + "total_token_count": 150 + } + } + """ + # Extract text from content.parts + text = None + content = chunk.get("content", {}) + parts = content.get("parts", []) + for part in parts: + if isinstance(part, dict) and "text" in part: + text = part["text"] + break + + # Extract finish_reason + finish_reason = None + raw_finish_reason = chunk.get("finish_reason") + if raw_finish_reason == "STOP": + finish_reason = "stop" + elif raw_finish_reason: + finish_reason = raw_finish_reason.lower() + + # Extract usage from usage_metadata + usage = None + usage_metadata = chunk.get("usage_metadata", {}) + if usage_metadata: + usage = ChatCompletionUsageBlock( + prompt_tokens=usage_metadata.get("prompt_token_count", 0), + completion_tokens=usage_metadata.get("candidates_token_count", 0), + total_tokens=usage_metadata.get("total_token_count", 0), + ) + + # Return ModelResponseStream (OpenAI-compatible chunk) + return ModelResponseStream( + choices=[ + StreamingChoices( + finish_reason=finish_reason, + index=0, + delta=Delta( + content=text, + role="assistant" if text else None, + ), + ) + ], + usage=usage, + ) diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py new file mode 100644 index 00000000000..4c07e8455e3 --- /dev/null +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -0,0 +1,508 @@ +""" +Transformation for Vertex AI Agent Engine (Reasoning Engines) + +Handles the transformation between LiteLLM's OpenAI-compatible format and +Vertex AI Reasoning Engine's API format. + +API Reference: +- :query endpoint - for session management (create, get, list, delete) +- :streamQuery endpoint - for actual queries (stream_query method) +""" + +import json +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast + +import httpx + +from litellm._logging import verbose_logger +from litellm._uuid import uuid +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, +) +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( + VertexAgentEngineResponseIterator, +) +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import Choices, Message, ModelResponse, Usage + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.utils import CustomStreamWrapper + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + HTTPHandler = Any + AsyncHTTPHandler = Any + CustomStreamWrapper = Any + + +class VertexAgentEngineError(BaseLLMException): + """Exception for Vertex Agent Engine errors.""" + + def __init__(self, status_code: int, message: str): + self.status_code = status_code + self.message = message + super().__init__(message=message, status_code=status_code) + + +class VertexAgentEngineConfig(BaseConfig, VertexBase): + """ + Configuration for Vertex AI Agent Engine (Reasoning Engines). + + Model format: vertex_ai/agent_engine/ + Where resource_id is the numeric ID of the reasoning engine. + """ + + def __init__(self, **kwargs): + BaseConfig.__init__(self, **kwargs) + VertexBase.__init__(self) + + def get_supported_openai_params(self, model: str) -> List[str]: + """Vertex Agent Engine has limited OpenAI compatible params.""" + return ["user"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """Map OpenAI params to Agent Engine params.""" + # Map 'user' to 'user_id' for session management + if "user" in non_default_params: + optional_params["user_id"] = non_default_params["user"] + return optional_params + + def _parse_model_string(self, model: str) -> Tuple[str, str]: + """ + Parse model string to extract resource ID. + + Model format: agent_engine/// + Or: agent_engine/ (uses default project/location) + + Returns: (resource_path, engine_id) + """ + # Remove 'agent_engine/' prefix if present + if model.startswith("agent_engine/"): + model = model[len("agent_engine/") :] + + # Check if it's a full resource path + if model.startswith("projects/"): + # Full path: projects/123/locations/us-central1/reasoningEngines/456 + return model, model.split("/")[-1] + + # Just the engine ID + return model, model + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete URL for the request. + + For Vertex Agent Engine: + - Non-streaming: :query endpoint (for session management) + - Streaming: :streamQuery endpoint (for actual queries) + """ + resource_path, engine_id = self._parse_model_string(model) + + # Get project and location from litellm_params or environment + vertex_project = self.safe_get_vertex_ai_project(litellm_params) + vertex_location = self.safe_get_vertex_ai_location(litellm_params) or "us-central1" + + # Build the full resource path if only engine_id was provided + if not resource_path.startswith("projects/"): + if not vertex_project: + raise ValueError( + "vertex_project is required for Vertex Agent Engine. " + "Set via litellm_params['vertex_project'] or VERTEXAI_PROJECT env var." + ) + resource_path = f"projects/{vertex_project}/locations/{vertex_location}/reasoningEngines/{engine_id}" + + # Build the base URL + base_url = f"https://{vertex_location}-aiplatform.googleapis.com" + + # Always use :streamQuery endpoint for actual queries + # The :query endpoint only supports session management methods + # (create_session, get_session, list_sessions, delete_session, etc.) + endpoint = f"{base_url}/v1beta1/{resource_path}:streamQuery" + + verbose_logger.debug(f"Vertex Agent Engine URL: {endpoint}") + return endpoint + + def _get_auth_headers( + self, + optional_params: dict, + litellm_params: dict, + ) -> Dict[str, str]: + """Get authentication headers using Google Cloud credentials.""" + vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) + + # Get access token using VertexBase + access_token, project_id = self.get_access_token( + credentials=vertex_credentials, + project_id=vertex_project, + ) + + verbose_logger.debug(f"Vertex Agent Engine: Authenticated for project {project_id}") + + return { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + } + + def _get_user_id(self, optional_params: dict) -> str: + """Get or generate user ID for session management.""" + user_id = optional_params.get("user_id") or optional_params.get("user") + if user_id: + return user_id + # Generate a user ID + return f"litellm-user-{str(uuid.uuid4())[:8]}" + + def _get_session_id(self, optional_params: dict) -> Optional[str]: + """Get session ID if provided.""" + return optional_params.get("session_id") + + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Transform the request to Vertex Agent Engine format. + + The API expects: + { + "class_method": "stream_query", + "input": { + "message": "...", + "user_id": "...", + "session_id": "..." (optional) + } + } + """ + # Use the last message content as the prompt + prompt = convert_content_list_to_str(messages[-1]) + + # Get user_id and session_id + user_id = self._get_user_id(optional_params) + session_id = self._get_session_id(optional_params) + + # Build the input + input_data: Dict[str, Any] = { + "message": prompt, + "user_id": user_id, + } + + if session_id: + input_data["session_id"] = session_id + + # Build the request payload + # Note: stream_query is used for both streaming and non-streaming + # The difference is the endpoint (:streamQuery vs :query) + payload = { + "class_method": "stream_query", + "input": input_data, + } + + verbose_logger.debug(f"Vertex Agent Engine payload: {payload}") + return payload + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """Validate environment and set up authentication headers.""" + auth_headers = self._get_auth_headers(optional_params, litellm_params) + headers.update(auth_headers) + return headers + + def _extract_text_from_response(self, response_data: dict) -> str: + """Extract text content from the response.""" + # Try to get from content.parts + content = response_data.get("content", {}) + parts = content.get("parts", []) + for part in parts: + if "text" in part: + return part["text"] + + # Try actions.state_delta + actions = response_data.get("actions", {}) + state_delta = actions.get("state_delta", {}) + for key, value in state_delta.items(): + if isinstance(value, str) and value: + return value + + return "" + + def _calculate_usage( + self, model: str, messages: List[AllMessageValues], content: str + ) -> Optional[Usage]: + """Calculate token usage using LiteLLM's token counter.""" + try: + from litellm.utils import token_counter + + prompt_tokens = token_counter(model="gpt-3.5-turbo", messages=messages) + completion_tokens = token_counter( + model="gpt-3.5-turbo", text=content, count_response_tokens=True + ) + total_tokens = prompt_tokens + completion_tokens + + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ) + except Exception as e: + verbose_logger.warning(f"Failed to calculate token usage: {str(e)}") + return None + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: LiteLLMLoggingObj, + request_data: dict, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + encoding: Any, + api_key: Optional[str] = None, + json_mode: Optional[bool] = None, + ) -> ModelResponse: + """ + Transform Vertex Agent Engine response to LiteLLM ModelResponse format. + + The response is a streaming SSE format even for non-streaming requests. + We need to collect all the chunks and extract the final response. + """ + try: + content_type = raw_response.headers.get("content-type", "").lower() + verbose_logger.debug(f"Vertex Agent Engine response Content-Type: {content_type}") + + # Parse the SSE response + response_text = raw_response.text + verbose_logger.debug(f"Response (first 500 chars): {response_text[:500]}") + + # Extract content from SSE stream + content = "" + for line in response_text.strip().split("\n"): + line = line.strip() + if not line: + continue + + try: + data = json.loads(line) + if isinstance(data, dict): + text = self._extract_text_from_response(data) + if text: + content = text # Use the last non-empty text + except json.JSONDecodeError: + continue + + # Create the message + message = Message(content=content, role="assistant") + + # Create choices + choice = Choices(finish_reason="stop", index=0, message=message) + + # Update model response + model_response.choices = [choice] + model_response.model = model + + # Calculate usage + calculated_usage = self._calculate_usage(model, messages, content) + if calculated_usage: + setattr(model_response, "usage", calculated_usage) + + return model_response + + except Exception as e: + verbose_logger.error(f"Error processing Vertex Agent Engine response: {str(e)}") + raise VertexAgentEngineError( + message=f"Error processing response: {str(e)}", + status_code=raw_response.status_code, + ) + + def get_streaming_response( + self, + model: str, + raw_response: httpx.Response, + ) -> VertexAgentEngineResponseIterator: + """Return a streaming iterator for SSE responses.""" + return VertexAgentEngineResponseIterator( + streaming_response=raw_response.iter_lines(), + sync_stream=True, + ) + + def get_sync_custom_stream_wrapper( + self, + model: str, + custom_llm_provider: str, + logging_obj: LiteLLMLoggingObj, + api_base: str, + headers: dict, + data: dict, + messages: list, + client: Optional[Union[HTTPHandler, "AsyncHTTPHandler"]] = None, + json_mode: Optional[bool] = None, + signed_json_body: Optional[bytes] = None, + ) -> "CustomStreamWrapper": + """Get a CustomStreamWrapper for synchronous streaming.""" + from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler, + _get_httpx_client, + ) + from litellm.utils import CustomStreamWrapper + + if client is None or not isinstance(client, HTTPHandler): + client = _get_httpx_client(params={}) + + # Avoid logging sensitive api_base directly + verbose_logger.debug("Making sync streaming request to Vertex AI endpoint.") + + # Make streaming request + response = client.post( + api_base, + headers=headers, + data=json.dumps(data), + stream=True, + logging_obj=logging_obj, + ) + + if response.status_code != 200: + raise VertexAgentEngineError( + status_code=response.status_code, message=str(response.read()) + ) + + # Create iterator for SSE stream + completion_stream = self.get_streaming_response(model=model, raw_response=response) + + streaming_response = CustomStreamWrapper( + completion_stream=completion_stream, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) + + # LOGGING + logging_obj.post_call( + input=messages, + api_key="", + original_response="first stream response received", + additional_args={"complete_input_dict": data}, + ) + + return streaming_response + + async def get_async_custom_stream_wrapper( + self, + model: str, + custom_llm_provider: str, + logging_obj: LiteLLMLoggingObj, + api_base: str, + headers: dict, + data: dict, + messages: list, + client: Optional["AsyncHTTPHandler"] = None, + json_mode: Optional[bool] = None, + signed_json_body: Optional[bytes] = None, + ) -> "CustomStreamWrapper": + """Get a CustomStreamWrapper for asynchronous streaming.""" + from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, + ) + from litellm.utils import CustomStreamWrapper + + if client is None or not isinstance(client, AsyncHTTPHandler): + client = get_async_httpx_client( + llm_provider=cast(Any, "vertex_ai"), params={} + ) + + # Avoid logging sensitive api_base directly + verbose_logger.debug("Making async streaming request to Vertex AI endpoint.") + + # Make async streaming request + response = await client.post( + api_base, + headers=headers, + data=json.dumps(data), + stream=True, + logging_obj=logging_obj, + ) + + if response.status_code != 200: + raise VertexAgentEngineError( + status_code=response.status_code, message=str(await response.aread()) + ) + + # Create iterator for SSE stream (async) + completion_stream = VertexAgentEngineResponseIterator( + streaming_response=response.aiter_lines(), + sync_stream=False, + ) + + streaming_response = CustomStreamWrapper( + completion_stream=completion_stream, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) + + # LOGGING + logging_obj.post_call( + input=messages, + api_key="", + original_response="first stream response received", + additional_args={"complete_input_dict": data}, + ) + + return streaming_response + + @property + def has_custom_stream_wrapper(self) -> bool: + """Indicates that this config has custom streaming support.""" + return True + + @property + def supports_stream_param_in_request_body(self) -> bool: + """Agent Engine does not allow passing `stream` in the request body.""" + return False + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> BaseLLMException: + return VertexAgentEngineError(status_code=status_code, message=error_message) + + def should_fake_stream( + self, + model: Optional[str], + stream: Optional[bool], + custom_llm_provider: Optional[str] = None, + ) -> bool: + """Agent Engine always returns SSE streams, so we use real streaming.""" + return False + diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 3cfa55c0606..6bb11430f20 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -5,7 +5,6 @@ from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, get_ty import httpx import litellm -from litellm.utils import supports_response_schema, supports_system_messages from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs @@ -14,6 +13,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import PartType, Schema from litellm.types.utils import TokenCountResponse +from litellm.utils import supports_response_schema, supports_system_messages class VertexAIError(BaseLLMException): @@ -36,6 +36,7 @@ class VertexAIModelRoute(str, Enum): MODEL_GARDEN = "model_garden" NON_GEMINI = "non_gemini" OPENAI_COMPATIBLE = "openai" + AGENT_ENGINE = "agent_engine" VERTEX_AI_MODEL_ROUTES = [f"{route.value}/" for route in VertexAIModelRoute] @@ -76,6 +77,10 @@ def get_vertex_ai_model_route( if litellm_params and litellm_params.get("base_model") is not None: if "gemini" in litellm_params["base_model"]: return VertexAIModelRoute.GEMINI + + # Check for agent_engine models (Reasoning Engines) + if "agent_engine/" in model: + return VertexAIModelRoute.AGENT_ENGINE # Check if numeric endpoint ID with custom api_base (PSC endpoint) # Route to GEMINI (HTTP path) to support PSC endpoints properly diff --git a/litellm/main.py b/litellm/main.py index b08ffd16e3d..4176c96d348 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3242,6 +3242,37 @@ def completion( # type: ignore # noqa: PLR0915 timeout=timeout, client=client, ) + elif model_route == VertexAIModelRoute.AGENT_ENGINE: + # Vertex AI Agent Engine (Reasoning Engines) + from litellm.llms.vertex_ai.agent_engine.transformation import ( + VertexAgentEngineConfig, + ) + + vertex_agent_engine_config = VertexAgentEngineConfig() + + # Update litellm_params with vertex credentials + litellm_params["vertex_project"] = vertex_ai_project + litellm_params["vertex_location"] = vertex_ai_location + litellm_params["vertex_credentials"] = vertex_credentials + + model_response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + model_response=model_response, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + encoding=encoding, + api_key=None, + api_base=api_base, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + client=client, + custom_llm_provider="vertex_ai", + provider_config=vertex_agent_engine_config, + headers=headers or {}, + ) else: # VertexAIModelRoute.NON_GEMINI model_response = vertex_ai_non_gemini.completion( model=model, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 01d7f076edc..2a7f8aa3ddf 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -30628,7 +30628,7 @@ "litellm_provider": "fireworks_ai", "mode": "embedding" }, - "fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": { + "fireworks_ai/accounts/fireworks/models/": { "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, diff --git a/litellm/proxy/public_endpoints/agent_create_fields.json b/litellm/proxy/public_endpoints/agent_create_fields.json index c559b61a76f..931c9a43498 100644 --- a/litellm/proxy/public_endpoints/agent_create_fields.json +++ b/litellm/proxy/public_endpoints/agent_create_fields.json @@ -166,6 +166,29 @@ "litellm_params_template": { "custom_llm_provider": "pydantic_ai_agents" } + }, + { + "agent_type": "vertex_agent_engine", + "agent_type_display_name": "Vertex AI Agent Engine", + "description": "Connect to Google Cloud Vertex AI Reasoning Engines", + "logo_url": "/ui/assets/logos/google.svg", + "inherit_credentials_from_provider": "Vertex_AI", + "model_template": "vertex_ai/agent_engine/{reasoning_engine_id}", + "credential_fields": [ + { + "key": "reasoning_engine_id", + "label": "Reasoning Engine Resource ID", + "placeholder": "projects/123456789/locations/us-central1/reasoningEngines/987654321", + "tooltip": "The full resource ID of your Vertex AI Reasoning Engine. Find this in Google Cloud Console under Vertex AI > Agent Builder > Your Agent.", + "required": true, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": false + } + ], + "litellm_params_template": { + "custom_llm_provider": "vertex_ai" + } } ] diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 629760a7dd2..68264a576fe 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2689,8 +2689,8 @@ "key": "vertex_credentials", "label": "Vertex Credentials", "placeholder": null, - "tooltip": null, - "required": true, + "tooltip": "Optional - Upload your GCP service account JSON file. If not provided, uses default GCP credentials (ADC).", + "required": false, "field_type": "upload", "options": null, "default_value": null diff --git a/tests/agent_tests/local_vertex_agent.py b/tests/agent_tests/local_vertex_agent.py new file mode 100644 index 00000000000..3cc9f868612 --- /dev/null +++ b/tests/agent_tests/local_vertex_agent.py @@ -0,0 +1,151 @@ +""" +Test script for Vertex AI Reasoning Engine. + +This script demonstrates how to: +1. Authenticate with Google Cloud +2. Send queries to a Vertex AI Reasoning Engine using the :query endpoint + +Usage: + python local_vertex_agent.py + +Requirements: + pip install httpx google-auth +""" + +import asyncio +import json +from uuid import uuid4 + +from google.auth import default +from google.auth.transport.requests import Request +import httpx + +# Configuration - update these for your agent +PROJECT_ID = "gen-lang-client-0682925754" # Your GCP project ID +LOCATION = "us-central1" # Your agent's location + +# For Reasoning Engines, use just the numeric ID at the end +REASONING_ENGINE_ID = "8263861224643493888" + +# The project number from the resource name +PROJECT_NUMBER = "1060139831167" + + +async def main(): + """Main function to test Vertex AI Reasoning Engine.""" + + # Step 1: Authenticate with Google Cloud + print("Step 1: Authenticating with Google Cloud...") + credentials, project = default(scopes=['https://www.googleapis.com/auth/cloud-platform']) + credentials.refresh(Request()) + print(f"Authenticated! Project: {project}") + print(f"Token (first 20 chars): {credentials.token[:20]}...") + + # Step 2: Build the endpoint URL + base_url = f"https://{LOCATION}-aiplatform.googleapis.com" + resource_path = f"projects/{PROJECT_NUMBER}/locations/{LOCATION}/reasoningEngines/{REASONING_ENGINE_ID}" + + # The Reasoning Engine uses :query endpoint with specific format + query_url = f"{base_url}/v1beta1/{resource_path}:query" + stream_url = f"{base_url}/v1beta1/{resource_path}:streamQuery" + + print(f"\nQuery URL: {query_url}") + print(f"Stream URL: {stream_url}") + + # Step 3: Create authenticated httpx client + print("\nStep 2: Creating authenticated HTTP client...") + client = httpx.AsyncClient( + headers={ + "Authorization": f"Bearer {credentials.token}", + "Content-Type": "application/json", + }, + timeout=120.0, + ) + + # Step 4: Build the query request (non-streaming) + # Note: For non-streaming, we need to: + # 1. Create a session + # 2. Use the streaming endpoint with stream_query method + # The :query endpoint only supports session management methods + + user_id = f"test-user-{uuid4().hex[:8]}" + + # First create a session + create_session_request = { + "class_method": "async_create_session", + "input": { + "user_id": user_id, + } + } + + print(f"\nStep 3: Creating session...") + print(f"User ID: {user_id}") + + async with client: + # Create session + print(f"\nSending to: {query_url}") + response = await client.post(query_url, json=create_session_request) + print(f"Create session status: {response.status_code}") + + if response.status_code == 200: + session_data = response.json() + print(f"Session created:\n{json.dumps(session_data, indent=2)}") + + # Extract session_id from response + session_id = session_data.get("output", {}).get("id") or session_data.get("output", {}).get("session_id") + print(f"\nSession ID: {session_id}") + + # Now send the actual query via streamQuery + query_request = { + "class_method": "stream_query", + "input": { + "message": "Hello! What can you do?", + "user_id": user_id, + "session_id": session_id, + } + } + + print(f"\nStep 4: Sending query via streamQuery...") + print(f"Request:\n{json.dumps(query_request, indent=2)}") + + # Use streaming endpoint but collect full response + async with client.stream("POST", stream_url, json=query_request) as stream_response: + print(f"Query status: {stream_response.status_code}") + + if stream_response.status_code == 200: + print("\nResponse:") + full_response = "" + async for line in stream_response.aiter_lines(): + if line: + full_response = line # Keep last line (full response) + + # Parse and display + try: + data = json.loads(full_response) + # Extract the text from the response + content = data.get("content", {}) + parts = content.get("parts", []) + for part in parts: + if "text" in part: + print(f"\nAgent response:\n{part['text']}") + except: + print(full_response) + else: + content = await stream_response.aread() + print(f"Error: {content.decode()}") + else: + print(f"Error creating session: {response.text}") + + +if __name__ == "__main__": + print("=" * 60) + print("Vertex AI Reasoning Engine Test Script") + print("=" * 60) + print(f"\nConfiguration:") + print(f" PROJECT_ID: {PROJECT_ID}") + print(f" PROJECT_NUMBER: {PROJECT_NUMBER}") + print(f" LOCATION: {LOCATION}") + print(f" REASONING_ENGINE_ID: {REASONING_ENGINE_ID}") + print() + + asyncio.run(main()) diff --git a/tests/agent_tests/test_a2a_completion_bridge.py b/tests/agent_tests/test_a2a_completion_bridge.py index 4191821f3de..224809dd7f5 100644 --- a/tests/agent_tests/test_a2a_completion_bridge.py +++ b/tests/agent_tests/test_a2a_completion_bridge.py @@ -201,3 +201,79 @@ async def test_a2a_completion_bridge_bedrock_agentcore(): print(f"Received {len(chunks)} chunks from Bedrock AgentCore") + +# ============================================================ +# Vertex AI Agent Engine Tests +# ============================================================ + +# Configuration - update these for your Vertex AI Reasoning Engine +VERTEX_AGENT_RESOURCE_NAME = "projects/1060139831167/locations/us-central1/reasoningEngines/8263861224643493888" + + +@pytest.mark.asyncio +async def test_vertex_agent_engine_non_streaming(): + """ + Test non-streaming request to Vertex AI Agent Engine via litellm.acompletion. + + Uses the Reasoning Engine resource ID to call a hosted agent. + """ + + litellm._turn_on_debug() + + # Call via litellm.acompletion with vertex_ai/agent_engine/ prefix + response = await litellm.acompletion( + model=f"vertex_ai/agent_engine/{VERTEX_AGENT_RESOURCE_NAME}", + messages=[{"role": "user", "content": "Hello! What can you do?"}], + stream=False, + ) + + print(f"\n=== Vertex Agent Engine Non-Streaming Response ===") + print(f"Response: {response}") + + # Basic assertions + assert response is not None + assert hasattr(response, "choices") + assert len(response.choices) > 0 + assert response.choices[0].message is not None + assert response.choices[0].message.content is not None + assert len(response.choices[0].message.content) > 0 + + print(f"Agent response: {response.choices[0].message.content[:200]}...") + + +@pytest.mark.asyncio +async def test_vertex_agent_engine_streaming(): + """ + Test streaming request to Vertex AI Agent Engine via litellm.acompletion. + + Uses the Reasoning Engine resource ID to call a hosted agent with streaming. + """ + #litellm._turn_on_debug() + + # Call via litellm.acompletion with streaming + response = await litellm.acompletion( + model=f"vertex_ai/agent_engine/{VERTEX_AGENT_RESOURCE_NAME}", + messages=[{"role": "user", "content": "Hello! What can you do?"}], + stream=True, + ) + + print(f"\n=== Vertex Agent Engine Streaming Response ===") + + chunks = [] + full_content = "" + async for chunk in response: + print(f"Chunk: {chunk}") + # chunks.append(chunk) + # if hasattr(chunk, "choices") and len(chunk.choices) > 0: + # delta = chunk.choices[0].delta + # if hasattr(delta, "content") and delta.content: + # full_content += delta.content + # print(f"Chunk: {delta.content}", end="", flush=True) + + # # print(f"\n\nReceived {len(chunks)} chunks") + # print(f"Full content: {full_content[:200]}...") + + # # Basic assertions + # assert len(chunks) > 0 + # assert len(full_content) > 0 + diff --git a/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py b/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py new file mode 100644 index 00000000000..cb3a5807d8c --- /dev/null +++ b/tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py @@ -0,0 +1,128 @@ +""" +Tests for Vertex AI Agent Engine transformation. + +Tests the request transformation and streaming chunk parsing without making real API calls. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( + VertexAgentEngineResponseIterator, +) +from litellm.llms.vertex_ai.agent_engine.transformation import VertexAgentEngineConfig + + +class TestVertexAgentEngineTransformRequest: + """Tests for transform_request method.""" + + def test_transform_request_basic(self): + """ + Test that transform_request correctly formats messages into Vertex Agent Engine payload. + """ + config = VertexAgentEngineConfig() + + messages = [{"role": "user", "content": "Hello, what can you do?"}] + optional_params = {"user_id": "test-user-123"} + litellm_params = {} + + result = config.transform_request( + model="agent_engine/123456789", + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers={}, + ) + + assert result["class_method"] == "stream_query" + assert result["input"]["message"] == "Hello, what can you do?" + assert result["input"]["user_id"] == "test-user-123" + assert "session_id" not in result["input"] + + def test_transform_request_with_session_id(self): + """ + Test that transform_request includes session_id when provided. + """ + config = VertexAgentEngineConfig() + + messages = [{"role": "user", "content": "Follow up question"}] + optional_params = { + "user_id": "test-user-123", + "session_id": "session-abc-456", + } + litellm_params = {} + + result = config.transform_request( + model="agent_engine/123456789", + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers={}, + ) + + assert result["class_method"] == "stream_query" + assert result["input"]["message"] == "Follow up question" + assert result["input"]["user_id"] == "test-user-123" + assert result["input"]["session_id"] == "session-abc-456" + + +class TestVertexAgentEngineChunkParser: + """Tests for the streaming chunk parser.""" + + def test_chunk_parser_with_text_content(self): + """ + Test that chunk_parser correctly extracts text from Vertex Agent Engine response format. + """ + iterator = VertexAgentEngineResponseIterator( + streaming_response=iter([]), + sync_stream=True, + ) + + chunk = { + "content": { + "parts": [{"text": "Hello! I can help you with financial analysis."}], + "role": "model", + }, + "finish_reason": "STOP", + "usage_metadata": { + "prompt_token_count": 100, + "candidates_token_count": 50, + "total_token_count": 150, + }, + } + + result = iterator.chunk_parser(chunk) + + assert result.choices[0].delta.content == "Hello! I can help you with financial analysis." + assert result.choices[0].delta.role == "assistant" + assert result.choices[0].finish_reason == "stop" + assert result.usage["prompt_tokens"] == 100 + assert result.usage["completion_tokens"] == 50 + assert result.usage["total_tokens"] == 150 + + def test_chunk_parser_without_finish_reason(self): + """ + Test that chunk_parser handles chunks without finish_reason (intermediate chunks). + """ + iterator = VertexAgentEngineResponseIterator( + streaming_response=iter([]), + sync_stream=True, + ) + + chunk = { + "content": { + "parts": [{"text": "Partial response..."}], + "role": "model", + }, + } + + result = iterator.chunk_parser(chunk) + + assert result.choices[0].delta.content == "Partial response..." + assert result.choices[0].finish_reason is None + assert result.usage is None + From 636efb7795a46ffbace42d94ff70142b1e1bca62 Mon Sep 17 00:00:00 2001 From: Igal Boxerman Date: Tue, 16 Dec 2025 05:19:03 +0200 Subject: [PATCH 031/121] feat(pillar): add masking support and MCP call support (#17959) - Add 'mask' action to SUPPORTED_ON_FLAGGED_ACTIONS - Automatically sanitizes sensitive content using masked_session_messages - Allows requests to proceed with masked content instead of blocking - Add MCP call support - Add pre_mcp_call and during_mcp_call to supported_event_hooks - Verify mcp_call is supported in call_type Literal types - Control exception details based on config - Conditionally include scanners/evidence in exceptions based on include_scanners and include_evidence settings - Reduces payload size when detailed exception info isn't needed - Add comprehensive test coverage - Tests for masking functionality - Tests for conditional exception details - Tests for MCP call support - Update documentation - Add Mask section explaining masking functionality - Clarify exception details control All changes maintain backward compatibility. --- .../docs/proxy/guardrails/pillar_security.md | 100 +++++- .../guardrail_hooks/pillar/pillar.py | 29 +- .../guardrails/test_pillar_guardrails.py | 299 ++++++++++++++++++ 3 files changed, 418 insertions(+), 10 deletions(-) diff --git a/docs/my-website/docs/proxy/guardrails/pillar_security.md b/docs/my-website/docs/proxy/guardrails/pillar_security.md index de0b0d53614..099919dc393 100644 --- a/docs/my-website/docs/proxy/guardrails/pillar_security.md +++ b/docs/my-website/docs/proxy/guardrails/pillar_security.md @@ -72,13 +72,15 @@ litellm --config config.yaml --port 4000 ### Overview -Pillar Security supports three execution modes for comprehensive protection: +Pillar Security supports five execution modes for comprehensive protection: | Mode | When It Runs | What It Protects | Use Case |------|-------------|------------------|---------- | **`pre_call`** | Before LLM call | User input only | Block malicious prompts, prevent prompt injection | **`during_call`** | Parallel with LLM call | User input only | Input monitoring with lower latency | **`post_call`** | After LLM response | Full conversation context | Output filtering, PII detection in responses +| **`pre_mcp_call`** | Before MCP tool call | MCP tool inputs | Validate and sanitize MCP tool call arguments +| **`during_mcp_call`** | During MCP tool call | MCP tool inputs | Real-time monitoring of MCP tool calls ### Why Dual Mode is Recommended @@ -198,6 +200,85 @@ litellm_settings: set_verbose: true # Enable detailed logging ``` + + + +**Best for:** +- 🔒 **PII Protection**: Automatically sanitize sensitive data before sending to LLM +- ✅ **Continue Workflows**: Allow requests to proceed with masked content +- 🛡️ **Zero Trust**: Never expose sensitive data to LLM models +- 📊 **Compliance**: Meet data privacy requirements without blocking legitimate requests + +```yaml +model_list: + - model_name: gpt-4.1-mini + litellm_params: + model: openai/gpt-4.1-mini + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "pillar-masking" + litellm_params: + guardrail: pillar + mode: "pre_call" # Scan input before LLM call + api_key: os.environ/PILLAR_API_KEY # Your Pillar API key + api_base: os.environ/PILLAR_API_BASE # Pillar API endpoint + on_flagged_action: "mask" # Mask sensitive content instead of blocking + persist_session: true # Keep records for investigation + include_scanners: true # Understand which scanners triggered + include_evidence: true # Capture evidence for analysis + default_on: true # Enable for all requests + +general_settings: + master_key: "YOUR_LITELLM_PROXY_MASTER_KEY" + +litellm_settings: + set_verbose: true +``` + +**How it works:** +1. User sends request with sensitive data: `"My email is john@example.com"` +2. Pillar detects PII and returns masked version: `"My email is [MASKED_EMAIL]"` +3. LiteLLM replaces original messages with masked messages +4. Request proceeds to LLM with sanitized content +5. User receives response without exposing sensitive data + + + + +**Best for:** +- 🤖 **Agent Workflows**: Protect MCP (Model Context Protocol) tool calls +- 🔒 **Tool Input Validation**: Scan arguments passed to MCP tools +- 🛡️ **Comprehensive Coverage**: Extend security to all LLM endpoints + +```yaml +model_list: + - model_name: gpt-4.1-mini + litellm_params: + model: openai/gpt-4.1-mini + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "pillar-mcp-guard" + litellm_params: + guardrail: pillar + mode: "pre_mcp_call" # Scan MCP tool call inputs + api_key: os.environ/PILLAR_API_KEY # Your Pillar API key + api_base: os.environ/PILLAR_API_BASE # Pillar API endpoint + on_flagged_action: "block" # Block malicious MCP calls + default_on: true # Enable for all MCP calls + +general_settings: + master_key: "YOUR_LITELLM_PROXY_MASTER_KEY" + +litellm_settings: + set_verbose: true +``` + +**MCP Modes:** +- `pre_mcp_call`: Scan MCP tool call inputs before execution +- `during_mcp_call`: Monitor MCP tool calls in real-time + @@ -251,6 +332,15 @@ Logs the violation but allows the request to proceed: on_flagged_action: "monitor" ``` +#### Mask +Automatically sanitizes sensitive content (PII, secrets, etc.) in your messages before sending them to the LLM: + +```yaml +on_flagged_action: "mask" +``` + +When masking is enabled, sensitive information is automatically replaced with masked versions, allowing requests to proceed safely without exposing sensitive data to the LLM. + **Response Headers:** You can opt in to receiving detection details in response headers by configuring `include_scanners: true` and/or `include_evidence: true`. When enabled, these headers are included for **every request**—not just flagged ones—enabling comprehensive metrics, false positive analysis, and threat investigation. @@ -383,7 +473,8 @@ export PILLAR_TIMEOUT="5.0" **Quick takeaways** - Every request still runs *all* Pillar scanners; these options only change what comes back. - Choose richer responses when you need audit trails, lighter responses when latency or cost matters. -- Blocking is controlled by LiteLLM’s `on_flagged_action` configuration—Pillar headers do not change block/monitor behaviour. +- Actions (block/monitor/mask) are controlled by LiteLLM's `on_flagged_action` configuration—Pillar headers are automatically set based on your config. +- When blocking (`on_flagged_action: "block"`), the `include_scanners` and `include_evidence` settings control what details are included in the exception response. Pillar Security executes the full scanner suite on each call. The settings below tune the Protect response headers LiteLLM sends, letting you balance fidelity, retention, and latency. @@ -415,9 +506,10 @@ include_evidence: true # → plr_evidence (default true in LiteLLM) ``` Use when you only care about whether Pillar detected a threat. - > **📝 Note:** `flagged: true` means Pillar’s scanners recommend blocking. Pillar only reports this verdict—LiteLLM enforces your policy via the `on_flagged_action` configuration (no Pillar header controls it): - > - `on_flagged_action: "block"` → LiteLLM raises a 400 guardrail error + > **📝 Note:** `flagged: true` means Pillar's scanners recommend blocking. Pillar only reports this verdict—LiteLLM enforces your policy via the `on_flagged_action` configuration: + > - `on_flagged_action: "block"` → LiteLLM raises a 400 guardrail error (exception includes scanners/evidence based on `include_scanners`/`include_evidence` settings) > - `on_flagged_action: "monitor"` → LiteLLM logs the threat but still returns the LLM response + > - `on_flagged_action: "mask"` → LiteLLM replaces messages with masked versions and allows the request to proceed - **Scanner breakdown** (`include_scanners=true`) ```json diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 0df610177e5..ef22b099300 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -164,7 +164,7 @@ class PillarGuardrail(CustomGuardrail): using the Pillar Security API. """ - SUPPORTED_ON_FLAGGED_ACTIONS = ["block", "monitor"] + SUPPORTED_ON_FLAGGED_ACTIONS = ["block", "monitor", "mask"] DEFAULT_ON_FLAGGED_ACTION = "monitor" SUPPORTED_FALLBACK_ACTIONS = ["allow", "block"] DEFAULT_FALLBACK_ACTION = "allow" @@ -280,6 +280,8 @@ class PillarGuardrail(CustomGuardrail): GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call, + GuardrailEventHooks.pre_mcp_call, + GuardrailEventHooks.during_mcp_call, ] super().__init__( @@ -773,6 +775,15 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.warning("Pillar Guardrail: Threat detected") if self.on_flagged_action == "block": self._raise_pillar_detection_exception(pillar_response) + elif self.on_flagged_action == "mask": + verbose_proxy_logger.info("Pillar Guardrail: Masking mode - masking flagged content") + masked_messages = pillar_response.get("masked_session_messages", []) + if masked_messages: + original_data["messages"] = masked_messages + else: + verbose_proxy_logger.warning( + "Pillar Guardrail: Masking requested but no masked_session_messages in response" + ) elif self.on_flagged_action == "monitor": verbose_proxy_logger.info("Pillar Guardrail: Monitoring mode - allowing flagged content to proceed") @@ -788,14 +799,20 @@ class PillarGuardrail(CustomGuardrail): Raises: HTTPException: Always raises with security detection details """ + pillar_response_dict = { + "session_id": pillar_response.get("session_id"), + } + + # Conditionally include scanners and evidence based on config + if self.include_scanners: + pillar_response_dict["scanners"] = pillar_response.get("scanners", {}) + if self.include_evidence: + pillar_response_dict["evidence"] = pillar_response.get("evidence", []) + error_detail = { "error": "Blocked by Pillar Security Guardrail", "detection_message": "Security threats detected", - "pillar_response": { - "session_id": pillar_response.get("session_id"), - "scanners": pillar_response.get("scanners", {}), - "evidence": pillar_response.get("evidence", []), - }, + "pillar_response": pillar_response_dict, } verbose_proxy_logger.warning("Pillar Guardrail: Request blocked - Security threats detected") diff --git a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py index 2e7443e889f..0607b0de981 100644 --- a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py @@ -1142,6 +1142,305 @@ def test_get_config_model(): assert hasattr(config_model, "ui_friendly_name") +# ============================================================================ +# MASKING TESTS +# ============================================================================ + + +@pytest.fixture +def pillar_masked_response(): + """Fixture providing a Pillar API response with masked messages.""" + return Response( + json={ + "session_id": "test-session-123", + "flagged": True, + "masked_session_messages": [ + {"role": "user", "content": "My email is [MASKED_EMAIL]"} + ], + "evidence": [ + { + "category": "pii", + "type": "email", + "evidence": "test@example.com", + } + ], + "scanners": { + "jailbreak": False, + "prompt_injection": False, + "pii": True, + "toxic_language": False, + }, + }, + status_code=200, + request=Request( + method="POST", url="https://api.pillar.security/api/v1/protect" + ), + ) + + +@pytest.fixture +def pillar_mask_guardrail(env_setup): + """Fixture providing a PillarGuardrail instance in mask mode.""" + return PillarGuardrail( + guardrail_name="pillar-mask", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + on_flagged_action="mask", + ) + + +@pytest.mark.asyncio +async def test_pre_call_hook_masking_mode( + pillar_mask_guardrail, + sample_request_data, + user_api_key_dict, + dual_cache, + pillar_masked_response, +): + """Test pre-call hook masks content when action is 'mask'.""" + original_messages = sample_request_data["messages"].copy() + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=pillar_masked_response, + ): + result = await pillar_mask_guardrail.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + # Messages should be replaced with masked messages + assert result["messages"] == pillar_masked_response.json()["masked_session_messages"] + assert result["messages"] != original_messages + + +@pytest.mark.asyncio +async def test_pre_call_hook_masking_no_masked_messages( + pillar_mask_guardrail, + sample_request_data, + user_api_key_dict, + dual_cache, +): + """Test masking mode when API doesn't return masked_session_messages.""" + response_no_mask = Response( + json={ + "session_id": "test-session-123", + "flagged": True, + # No masked_session_messages + }, + status_code=200, + request=Request( + method="POST", url="https://api.pillar.security/api/v1/protect" + ), + ) + + original_messages = sample_request_data["messages"].copy() + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=response_no_mask, + ): + result = await pillar_mask_guardrail.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + # Messages should remain unchanged if no masked messages provided + assert result["messages"] == original_messages + + +# ============================================================================ +# CONDITIONAL EXCEPTION DETAILS TESTS +# ============================================================================ + + +@pytest.mark.asyncio +async def test_exception_without_scanners( + sample_request_data, + user_api_key_dict, + dual_cache, + pillar_flagged_response, +): + """Test exception excludes scanners when include_scanners is False.""" + guardrail = PillarGuardrail( + guardrail_name="pillar-no-scanners", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + on_flagged_action="block", + include_scanners=False, + include_evidence=True, + ) + + with pytest.raises(HTTPException) as excinfo: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=pillar_flagged_response, + ): + await guardrail.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + error_detail = excinfo.value.detail + assert "pillar_response" in error_detail + assert "scanners" not in error_detail["pillar_response"] + assert "evidence" in error_detail["pillar_response"] + + +@pytest.mark.asyncio +async def test_exception_without_evidence( + sample_request_data, + user_api_key_dict, + dual_cache, + pillar_flagged_response, +): + """Test exception excludes evidence when include_evidence is False.""" + guardrail = PillarGuardrail( + guardrail_name="pillar-no-evidence", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + on_flagged_action="block", + include_scanners=True, + include_evidence=False, + ) + + with pytest.raises(HTTPException) as excinfo: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=pillar_flagged_response, + ): + await guardrail.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + error_detail = excinfo.value.detail + assert "pillar_response" in error_detail + assert "scanners" in error_detail["pillar_response"] + assert "evidence" not in error_detail["pillar_response"] + + +@pytest.mark.asyncio +async def test_exception_without_scanners_or_evidence( + sample_request_data, + user_api_key_dict, + dual_cache, + pillar_flagged_response, +): + """Test exception excludes both scanners and evidence when both are False.""" + guardrail = PillarGuardrail( + guardrail_name="pillar-minimal", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + on_flagged_action="block", + include_scanners=False, + include_evidence=False, + ) + + with pytest.raises(HTTPException) as excinfo: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=pillar_flagged_response, + ): + await guardrail.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + error_detail = excinfo.value.detail + assert "pillar_response" in error_detail + pillar_response = error_detail["pillar_response"] + assert "scanners" not in pillar_response + assert "evidence" not in pillar_response + assert "session_id" in pillar_response # session_id should always be present + + +# ============================================================================ +# MCP CALL SUPPORT TESTS +# ============================================================================ + + +@pytest.mark.asyncio +async def test_pre_call_hook_mcp_call( + pillar_guardrail_instance, + sample_request_data, + user_api_key_dict, + dual_cache, + pillar_clean_response, +): + """Test pre-call hook works with MCP call type.""" + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=pillar_clean_response, + ): + result = await pillar_guardrail_instance.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="mcp_call", + ) + + assert result == sample_request_data + + +@pytest.mark.asyncio +async def test_moderation_hook_mcp_call( + pillar_guardrail_instance, + sample_request_data, + user_api_key_dict, + pillar_clean_response, +): + """Test moderation hook works with MCP call type.""" + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=pillar_clean_response, + ): + result = await pillar_guardrail_instance.async_moderation_hook( + data=sample_request_data, + user_api_key_dict=user_api_key_dict, + call_type="mcp_call", + ) + + assert result == sample_request_data + + +@pytest.mark.asyncio +async def test_mcp_call_masking( + pillar_mask_guardrail, + sample_request_data, + user_api_key_dict, + dual_cache, + pillar_masked_response, +): + """Test masking works with MCP call type.""" + original_messages = sample_request_data["messages"].copy() + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=pillar_masked_response, + ): + result = await pillar_mask_guardrail.async_pre_call_hook( + data=sample_request_data, + cache=dual_cache, + user_api_key_dict=user_api_key_dict, + call_type="mcp_call", + ) + + # Messages should be replaced with masked messages + assert result["messages"] == pillar_masked_response.json()["masked_session_messages"] + assert result["messages"] != original_messages + + if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) From f90dea7315c424fc3468c03b05dcf6e1baf12695 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 16 Dec 2025 08:50:21 +0530 Subject: [PATCH 032/121] fix(docker-compose.yml): move to docker.litellm.ai --- docker-compose.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker-compose.yml b/docker-compose.yml index 8898aff62da..988860a7877 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -4,7 +4,7 @@ services: context: . args: target: runtime - image: ghcr.io/berriai/litellm:main-stable + image: docker.litellm.ai/berriai/litellm:main-stable ######################################### ## Uncomment these lines to start proxy with a config.yaml file ## # volumes: From 0eb7d975a17fa62ff73a0b1abd9de608b78b049e Mon Sep 17 00:00:00 2001 From: Xingjian Li <43537913+OlivverX@users.noreply.github.com> Date: Tue, 16 Dec 2025 11:33:40 +0800 Subject: [PATCH 033/121] fix: Support signed URLs with query parameters for Vertex AI Gemini (#17976) - Fix _get_image_mime_type_from_url() to parse path without query params - Resolves issues with Tencent Cloud COS and other signed URLs --- .../litellm_core_utils/prompt_templates/common_utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index d2c91f4a841..ca2a092dbc8 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -689,7 +689,14 @@ def _get_image_mime_type_from_url(url: str) -> Optional[str]: video/mpegps video/flv """ + from urllib.parse import urlparse + url = url.lower() + + # Parse URL to extract path without query parameters + # This handles URLs like: https://example.com/image.jpg?signature=... + parsed = urlparse(url) + path = parsed.path # Map file extensions to mime types mime_types = { @@ -717,7 +724,7 @@ def _get_image_mime_type_from_url(url: str) -> Optional[str]: # Check each extension group against the URL for extensions, mime_type in mime_types.items(): - if any(url.endswith(ext) for ext in extensions): + if any(path.endswith(ext) for ext in extensions): return mime_type return None From ce113f4e4b61160d2c06e2ea7d53fa3121fbb683 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 15 Dec 2025 19:43:12 -0800 Subject: [PATCH 034/121] [Docs] Add docs on using pydantic ai agents with LiteLLM A2a gateway (#18026) * init A2AProviderConfigManager * move file * move file * add pydnatic ai folder * init providers * test_pydantic_ai_non_streaming * fix import * INIT pydantic * use_a2a_form_fields * test_vertex_agent_engine_streaming * add agent_engine * init transform for agent engine * init agent engine * VertexAgentEngineSSEStreamIterator * sample * ui add new fields * fix vertex_credentials * working SSE iterator * TestVertexAgentEngineTransformRequest * fix code QA check * stash docs * docs fix * fix logo * docs fix * doc pydantic ai * docs pydantic ai * new provider * docs fix --- docs/my-website/docs/a2a.md | 7 +- .../docs/providers/pydantic_ai_agent.md | 121 ++++++++++++++++++ docs/my-website/sidebars.js | 1 + .../out/assets/logos/pydantic.svg | 5 + provider_endpoints_support.json | 17 +++ 5 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 docs/my-website/docs/providers/pydantic_ai_agent.md create mode 100644 litellm/proxy/_experimental/out/assets/logos/pydantic.svg diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md index 9c94a2fbf29..b41daa37bfb 100644 --- a/docs/my-website/docs/a2a.md +++ b/docs/my-website/docs/a2a.md @@ -16,7 +16,7 @@ Add A2A Agents on LiteLLM AI Gateway, Invoke agents in A2A Protocol, track reque | Feature | Supported | |---------|-----------| -| Supported Agent Providers | A2A, LangGraph, Azure AI Foundry, Bedrock AgentCore | +| Supported Agent Providers | A2A, Pydantic AI, LangGraph, Azure AI Foundry, Bedrock AgentCore | | Logging | ✅ | | Load Balancing | ✅ | | Streaming | ✅ | @@ -45,6 +45,7 @@ You can add A2A-compatible agents through the LiteLLM Admin UI. The URL should be the invocation URL for your A2A agent (e.g., `http://localhost:10001`). + ### Add Azure AI Foundry Agents Follow [this guide, to add your azure ai foundry agent to LiteLLM Agent Gateway](./providers/azure_ai_agents#litellm-a2a-gateway) @@ -57,6 +58,10 @@ Follow [this guide, to add your langgraph agent to LiteLLM Agent Gateway](./prov Follow [this guide, to add your bedrock agentcore agent to LiteLLM Agent Gateway](./providers/bedrock_agentcore#litellm-a2a-gateway) +### Add Pydantic AI Agents + +Follow [this guide, to add your pydantic ai agent to LiteLLM Agent Gateway](./providers/pydantic_ai_agent#litellm-a2a-gateway) + ## Invoking your Agents Use the [A2A Python SDK](https://pypi.org/project/a2a/) to invoke agents through LiteLLM. diff --git a/docs/my-website/docs/providers/pydantic_ai_agent.md b/docs/my-website/docs/providers/pydantic_ai_agent.md new file mode 100644 index 00000000000..e96295faaf3 --- /dev/null +++ b/docs/my-website/docs/providers/pydantic_ai_agent.md @@ -0,0 +1,121 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Pydantic AI Agents + +Call Pydantic AI Agents via LiteLLM's A2A Gateway. + +| Property | Details | +|----------|---------| +| Description | Pydantic AI agents with native A2A support via the `to_a2a()` method. LiteLLM provides fake streaming support for agents that don't natively stream. | +| Provider Route on LiteLLM | A2A Gateway | +| Supported Endpoints | `/v1/a2a/message/send` | +| Provider Doc | [Pydantic AI Agents ↗](https://ai.pydantic.dev/agents/) | + +## LiteLLM A2A Gateway + +All Pydantic AI agents need to be exposed as A2A agents using the `to_a2a()` method. Once your agent server is running, you can add it to the LiteLLM Gateway. + +### 1. Setup Pydantic AI Agent Server + +LiteLLM requires Pydantic AI agents to follow the [A2A (Agent-to-Agent) protocol](https://github.com/google/A2A). Pydantic AI has native A2A support via the `to_a2a()` method, which exposes your agent as an A2A-compliant server. + +#### Install Dependencies + +```bash +pip install pydantic-ai fasta2a uvicorn +``` + +#### Create Agent + +```python title="agent.py" +from pydantic_ai import Agent + +agent = Agent('openai:gpt-4o-mini', instructions='Be helpful!') + +@agent.tool_plain +def get_weather(city: str) -> str: + """Get weather for a city.""" + return f"Weather in {city}: Sunny, 72°F" + +@agent.tool_plain +def calculator(expression: str) -> str: + """Evaluate a math expression.""" + return str(eval(expression)) + +# Native A2A server - Pydantic AI handles it automatically +app = agent.to_a2a() +``` + +#### Run Server + +```bash +uvicorn agent:app --host 0.0.0.0 --port 9999 +``` + +Server runs at `http://localhost:9999` + +### 2. Navigate to Agents + +From the sidebar, click "Agents" to open the agent management page, then click "+ Add New Agent". + +### 3. Select Pydantic AI Agent Type + +Click "A2A Standard" to see available agent types, then select "Pydantic AI". + +![Select A2A Standard](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/1055acb1-064b-4465-8e6a-8278291bc661/ascreenshot.jpeg?tl_px=0,0&br_px=2201,1230&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=395,147) + +![Select Pydantic AI](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/0998e38c-8534-40f1-931a-be96c2cae0ad/ascreenshot.jpeg?tl_px=0,52&br_px=2201,1283&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=421,277) + +### 4. Configure the Agent + +Fill in the following fields: + +- **Agent Name** - A unique identifier for your agent (e.g., `test-pydantic-agent`) +- **Agent URL** - The URL where your Pydantic AI agent is running. We use `http://localhost:9999` because that's where we started our Pydantic AI agent server in the previous step. + +![Enter Agent Name](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/8cf3fbde-05f3-48d1-81b6-6f857bd6d360/ascreenshot.jpeg?tl_px=0,0&br_px=2201,1230&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=443,225) + +![Configure Agent Name](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/fb555808-4761-4c49-a415-200ac1bdb525/ascreenshot.jpeg?tl_px=0,0&br_px=2617,1463&force_format=jpeg&q=100&width=1120.0) + +![Enter Agent URL](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/303eae61-4352-4fb0-a537-806839c234ba/ascreenshot.jpeg?tl_px=0,212&br_px=2201,1443&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=456,277) + +### 5. Create Agent + +Click "Create Agent" to save your configuration. + +![Create Agent](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/914f3367-df7d-4244-bd4d-e99ce0a6193a/ascreenshot.jpeg?tl_px=416,438&br_px=2618,1669&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=690,277) + +### 6. Test in Playground + +Go to "Playground" in the sidebar to test your agent. + +![Go to Playground](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/c73c9f3b-22af-4105-aafa-2d34c4986ef3/ascreenshot.jpeg?tl_px=0,0&br_px=2201,1230&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=44,97) + +### 7. Select A2A Endpoint + +Click the endpoint dropdown and search for "a2a", then select `/v1/a2a/message/send`. + +![Click Endpoint Dropdown](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/196d97ac-bcba-47f0-9880-97b80250e00c/ascreenshot.jpeg?tl_px=0,0&br_px=2201,1230&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=261,230) + +![Search for A2A](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/26b68f21-29f9-4c4c-b8b5-d2e11cbfd14a/ascreenshot.jpeg?tl_px=0,0&br_px=2617,1463&force_format=jpeg&q=100&width=1120.0) + +![Select A2A Endpoint](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/41576fb1-d385-4fb2-84e9-142dd7fe5181/ascreenshot.jpeg?tl_px=0,0&br_px=2201,1230&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=307,270) + +### 8. Select Your Agent and Send a Message + +Pick your Pydantic AI agent from the dropdown and send a test message. + +![Click Agent Dropdown](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/a96d7967-3d54-4cbf-bd3e-b38f1be9df76/ascreenshot.jpeg?tl_px=0,54&br_px=2201,1285&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=274,277) + +![Select Agent](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/e05a5a6e-d044-4480-b94e-7c03cfb92ac5/ascreenshot.jpeg?tl_px=0,113&br_px=2201,1344&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=290,277) + +![Send Message](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/29162702-968a-401a-aac1-c844bfc5f4a3/ascreenshot.jpeg?tl_px=91,653&br_px=2292,1883&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=524,436) + + +## Further Reading + +- [Pydantic AI Documentation](https://ai.pydantic.dev/) +- [Pydantic AI Agents](https://ai.pydantic.dev/agents/) +- [A2A Agent Gateway](../a2a.md) +- [A2A Cost Tracking](../a2a_cost_tracking.md) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 954a27d3182..3c8b23f856d 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -738,6 +738,7 @@ const sidebars = { "providers/petals", "providers/publicai", "providers/predibase", + "providers/pydantic_ai_agent", "providers/ragflow", "providers/recraft", "providers/replicate", diff --git a/litellm/proxy/_experimental/out/assets/logos/pydantic.svg b/litellm/proxy/_experimental/out/assets/logos/pydantic.svg new file mode 100644 index 00000000000..0ff8e5c44c7 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/pydantic.svg @@ -0,0 +1,5 @@ + + + diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index cec9f9e37c0..c5da3052738 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1946,6 +1946,23 @@ "rerank": false, "a2a": true } + }, + "pydantic_ai_agents": { + "display_name": "Pydantic AI Agents (`pydantic_ai_agents`)", + "url": "https://docs.litellm.ai/docs/providers/pydantic_ai_agent", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true + } } } } \ No newline at end of file From e40ad5203eb3a6a07a7d171f60f642f97ac75666 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 16 Dec 2025 09:28:04 +0530 Subject: [PATCH 035/121] Add container field as provider specific field --- litellm/llms/anthropic/chat/transformation.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 628121ab11c..72c95eef12b 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1132,22 +1132,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if content["type"] == "text": text_content += content["text"] ## TOOL CALLING - elif content["type"] == "tool_use": + elif content["type"] == "tool_use" or content["type"] == "server_tool_use": tool_call = AnthropicConfig.convert_tool_use_to_openai_format( anthropic_tool_content=content, index=idx, ) tool_calls.append(tool_call) - ## SERVER TOOL USE (for tool search) - elif content["type"] == "server_tool_use": - # Server tool use blocks are for tool search - treat as tool calls - # Note: using .get("input", {}) for server_tool_use as input may not be present - content_with_input = {**content, "input": content.get("input", {})} - tool_call = AnthropicConfig.convert_tool_use_to_openai_format( - anthropic_tool_content=content_with_input, - index=idx, - ) - tool_calls.append(tool_call) ## TOOL SEARCH TOOL RESULT (skip - this is metadata about tool discovery) elif content["type"] == "tool_search_tool_result": # This block contains tool_references that were discovered @@ -1343,6 +1333,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "context_management" ) + container: Optional[Dict] = completion_response.get("container") + provider_specific_fields: Dict[str, Any] = { "citations": citations, "thinking_blocks": thinking_blocks, @@ -1351,7 +1343,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): provider_specific_fields["context_management"] = context_management if web_search_results is not None: provider_specific_fields["web_search_results"] = web_search_results - + if container is not None: + provider_specific_fields["container"] = container + _message = litellm.Message( tool_calls=tool_calls, content=text_content or None, From cd9b71093ac6c19f68ae5df214e9d1709e24205b Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 16 Dec 2025 01:08:39 -0300 Subject: [PATCH 036/121] Improve issue labeling: add component dropdown and more provider keywords (#17957) - Replace "ML Ops Team" dropdown with "What part of LiteLLM is this about?" with options: SDK, Proxy, UI Dashboard, Docs, Other - Add component dropdown to feature_request.yml template - Rename label-mlops.yml to label-component.yml and update to auto-label issues with sdk, proxy, ui-dashboard, or docs based on selection - Add more provider keywords to issue-keyword-labeler: gemini, cohere, mistral, groq, ollama, deepseek --- .github/ISSUE_TEMPLATE/bug_report.yml | 12 +- .github/ISSUE_TEMPLATE/feature_request.yml | 12 ++ .github/workflows/issue-keyword-labeler.yml | 2 +- .github/workflows/label-component.yml | 144 ++++++++++++++++++++ .github/workflows/label-mlops.yml | 17 --- 5 files changed, 164 insertions(+), 23 deletions(-) create mode 100644 .github/workflows/label-component.yml delete mode 100644 .github/workflows/label-mlops.yml diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 8fbf1b3c5b4..39b46cba999 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -23,13 +23,15 @@ body: description: Please copy and paste any relevant log output. This will be automatically formatted into code, so no need for backticks. render: shell - type: dropdown - id: ml-ops-team + id: component attributes: - label: Are you a ML Ops Team? - description: This helps us prioritize your requests correctly + label: What part of LiteLLM is this about? options: - - "No" - - "Yes" + - "SDK (litellm Python package)" + - "Proxy" + - "UI Dashboard" + - "Docs" + - "Other" validations: required: true - type: input diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index 13a2132ec95..96b95cc7f02 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -22,6 +22,18 @@ body: description: Please outline the motivation for the proposal. Is your feature request related to a specific problem? e.g., "I'm working on X and would like Y to be possible". If this is related to another GitHub issue, please link here too. validations: required: true + - type: dropdown + id: component + attributes: + label: What part of LiteLLM is this about? + options: + - "SDK (litellm Python package)" + - "Proxy" + - "UI Dashboard" + - "Docs" + - "Other" + validations: + required: true - type: dropdown id: hiring-interest attributes: diff --git a/.github/workflows/issue-keyword-labeler.yml b/.github/workflows/issue-keyword-labeler.yml index 60c18e3b9af..936f90f747f 100644 --- a/.github/workflows/issue-keyword-labeler.yml +++ b/.github/workflows/issue-keyword-labeler.yml @@ -19,7 +19,7 @@ jobs: id: scan env: PROVIDER_ISSUE_WEBHOOK_URL: ${{ secrets.PROVIDER_ISSUE_WEBHOOK_URL }} - KEYWORDS: azure,openai,bedrock,vertexai,vertex ai,anthropic + KEYWORDS: azure,openai,bedrock,vertexai,vertex ai,anthropic,gemini,cohere,mistral,groq,ollama,deepseek run: python3 .github/scripts/scan_keywords.py - name: Ensure label exists diff --git a/.github/workflows/label-component.yml b/.github/workflows/label-component.yml new file mode 100644 index 00000000000..c0f9436288c --- /dev/null +++ b/.github/workflows/label-component.yml @@ -0,0 +1,144 @@ +name: Label Component Issues + +on: + issues: + types: + - opened + +jobs: + add-component-label: + runs-on: ubuntu-latest + permissions: + issues: write + steps: + - name: Add SDK label + if: contains(github.event.issue.body, 'SDK (litellm Python package)') + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const labelName = 'sdk'; + try { + await github.rest.issues.getLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName + }); + } catch (error) { + if (error.status === 404) { + await github.rest.issues.createLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName, + color: '0E7C86', + description: 'Issues related to the litellm Python SDK' + }); + } else { + throw error; + } + } + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + labels: [labelName] + }); + + - name: Add Proxy label + if: contains(github.event.issue.body, 'Proxy') + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const labelName = 'proxy'; + try { + await github.rest.issues.getLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName + }); + } catch (error) { + if (error.status === 404) { + await github.rest.issues.createLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName, + color: '5319E7', + description: 'Issues related to the LiteLLM Proxy' + }); + } else { + throw error; + } + } + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + labels: [labelName] + }); + + - name: Add UI Dashboard label + if: contains(github.event.issue.body, 'UI Dashboard') + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const labelName = 'ui-dashboard'; + try { + await github.rest.issues.getLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName + }); + } catch (error) { + if (error.status === 404) { + await github.rest.issues.createLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName, + color: 'D876E3', + description: 'Issues related to the LiteLLM UI Dashboard' + }); + } else { + throw error; + } + } + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + labels: [labelName] + }); + + - name: Add Docs label + if: contains(github.event.issue.body, 'Docs') + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const labelName = 'docs'; + try { + await github.rest.issues.getLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName + }); + } catch (error) { + if (error.status === 404) { + await github.rest.issues.createLabel({ + owner: context.repo.owner, + repo: context.repo.repo, + name: labelName, + color: 'FBCA04', + description: 'Issues related to LiteLLM documentation' + }); + } else { + throw error; + } + } + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + labels: [labelName] + }); diff --git a/.github/workflows/label-mlops.yml b/.github/workflows/label-mlops.yml deleted file mode 100644 index 37789c1ea76..00000000000 --- a/.github/workflows/label-mlops.yml +++ /dev/null @@ -1,17 +0,0 @@ -name: Label ML Ops Team Issues - -on: - issues: - types: - - opened - -jobs: - add-mlops-label: - runs-on: ubuntu-latest - steps: - - name: Check if ML Ops Team is selected - uses: actions-ecosystem/action-add-labels@v1 - if: contains(github.event.issue.body, '### Are you a ML Ops Team?') && contains(github.event.issue.body, 'Yes') - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - labels: "mlops user request" From 13c0ab898595d8cffe3acba706c91dfbeed3a9a8 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 16 Dec 2025 01:09:22 -0300 Subject: [PATCH 037/121] Cleanup PR template: remove redundant fields (#17956) - Remove Title section (already in PR title) - Remove screenshot requirement (CI validates tests) Co-authored-by: Krish Dholakia --- .github/pull_request_template.md | 5 ----- 1 file changed, 5 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 8977332ee01..b91b16c955c 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,7 +1,3 @@ -## Title - - - ## Relevant issues @@ -11,7 +7,6 @@ **Please complete all items before asking a LiteLLM maintainer to review your PR** - [ ] I have Added testing in the [`tests/litellm/`](https://github.com/BerriAI/litellm/tree/main/tests/litellm) directory, **Adding at least 1 test is a hard requirement** - [see details](https://docs.litellm.ai/docs/extras/contributing_code) -- [ ] I have added a screenshot of my new test passing locally - [ ] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code) - [ ] My PR's scope is as isolated as possible, it only solves 1 specific problem From 50606bf09026229b3ef9cf8ac98f6329ea96ef24 Mon Sep 17 00:00:00 2001 From: Eric84626 <97266539+Eric84626@users.noreply.github.com> Date: Tue, 16 Dec 2025 12:10:22 +0800 Subject: [PATCH 038/121] Added new step into rotate master key function for processing credentials table (#17952) * fix: Return 403 exception when calling GET responses api * fix: added new step into rotate master key function for processing credentials table --- .../proxy/credential_endpoints/endpoints.py | 9 ++--- .../key_management_endpoints.py | 34 +++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 647abb73648..9f228bb1184 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -21,11 +21,11 @@ router = APIRouter() class CredentialHelperUtils: @staticmethod - def encrypt_credential_values(credential: CredentialItem) -> CredentialItem: + def encrypt_credential_values(credential: CredentialItem, new_encryption_key: Optional[str] = None) -> CredentialItem: """Encrypt values in credential.credential_values and add to DB""" encrypted_credential_values = {} for key, value in (credential.credential_values or {}).items(): - encrypted_credential_values[key] = encrypt_value_helper(value) + encrypted_credential_values[key] = encrypt_value_helper(value, new_encryption_key) # Return a new object to avoid mutating the caller's credential, which # is kept in memory and should remain unencrypted. @@ -246,7 +246,7 @@ async def delete_credential( def update_db_credential( - db_credential: CredentialItem, updated_patch: CredentialItem + db_credential: CredentialItem, updated_patch: CredentialItem, new_encryption_key: Optional[str] = None ) -> CredentialItem: """ Update a credential in the DB. @@ -258,7 +258,8 @@ def update_db_credential( ) encrypted_credential = CredentialHelperUtils.encrypt_credential_values( - updated_patch + updated_patch, + new_encryption_key, ) # update model name if encrypted_credential.credential_name: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index da44bda791d..8ea3122ce01 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2539,6 +2539,40 @@ async def _rotate_master_key( new_master_key=new_master_key, ) + # 5. process credentials table + try: + credentials = await prisma_client.db.litellm_credentialstable.find_many() + except Exception: + credentials = None + if credentials: + from litellm.proxy.credential_endpoints.endpoints import update_db_credential + + for cred in credentials: + try: + decrypted_cred = proxy_config.decrypt_credentials(cred) + encrypted_cred = update_db_credential( + db_credential=cred, + updated_patch=decrypted_cred, + new_encryption_key=new_master_key, + ) + credential_object_jsonified = jsonify_object(encrypted_cred.model_dump()) + await prisma_client.db.litellm_credentialstable.update( + where={"credential_name": cred.credential_name}, + data={ + **credential_object_jsonified, + "updated_by": user_api_key_dict.user_id, + }, + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to re-encrypt credential {cred.credential_name}: {str(e)}" + ) + # Continue with next credential instead of failing entire rotation + continue + verbose_proxy_logger.debug( + f"Successfully re-encrypted {len(credentials)} credentials with new master key" + ) + def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: if data and data.new_key is not None: From 244d83ff474f129a66a5b645b4d032f6fa61df17 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 15 Dec 2025 20:11:40 -0800 Subject: [PATCH 039/121] [Docs] Litellm add docs vertex ai engine (#18027) * new provider doc * add to sidebar * stash docs * docs fix * docs vertex agent engine --- docs/my-website/docs/a2a.md | 10 +- .../docs/providers/vertex_ai_agent_engine.md | 216 ++++++++++++++++++ docs/my-website/sidebars.js | 1 + provider_endpoints_support.json | 17 ++ 4 files changed, 241 insertions(+), 3 deletions(-) create mode 100644 docs/my-website/docs/providers/vertex_ai_agent_engine.md diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md index b41daa37bfb..d7145e4b83c 100644 --- a/docs/my-website/docs/a2a.md +++ b/docs/my-website/docs/a2a.md @@ -16,7 +16,7 @@ Add A2A Agents on LiteLLM AI Gateway, Invoke agents in A2A Protocol, track reque | Feature | Supported | |---------|-----------| -| Supported Agent Providers | A2A, Pydantic AI, LangGraph, Azure AI Foundry, Bedrock AgentCore | +| Supported Agent Providers | A2A, Vertex AI Agent Engine, LangGraph, Azure AI Foundry, Bedrock AgentCore, Pydantic AI | | Logging | ✅ | | Load Balancing | ✅ | | Streaming | ✅ | @@ -50,14 +50,18 @@ The URL should be the invocation URL for your A2A agent (e.g., `http://localhost Follow [this guide, to add your azure ai foundry agent to LiteLLM Agent Gateway](./providers/azure_ai_agents#litellm-a2a-gateway) -### Add LangGraph Agents +### Add Vertex AI Agent Engine -Follow [this guide, to add your langgraph agent to LiteLLM Agent Gateway](./providers/langgraph#litellm-a2a-gateway) +Follow [this guide, to add your Vertex AI Agent Engine to LiteLLM Agent Gateway](./providers/vertex_ai_agent_engine) ### Add Bedrock AgentCore Agents Follow [this guide, to add your bedrock agentcore agent to LiteLLM Agent Gateway](./providers/bedrock_agentcore#litellm-a2a-gateway) +### Add LangGraph Agents + +Follow [this guide, to add your langgraph agent to LiteLLM Agent Gateway](./providers/langgraph#litellm-a2a-gateway) + ### Add Pydantic AI Agents Follow [this guide, to add your pydantic ai agent to LiteLLM Agent Gateway](./providers/pydantic_ai_agent#litellm-a2a-gateway) diff --git a/docs/my-website/docs/providers/vertex_ai_agent_engine.md b/docs/my-website/docs/providers/vertex_ai_agent_engine.md new file mode 100644 index 00000000000..3bd40e98684 --- /dev/null +++ b/docs/my-website/docs/providers/vertex_ai_agent_engine.md @@ -0,0 +1,216 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Vertex AI Agent Engine + +Call Vertex AI Agent Engine (Reasoning Engines) in the OpenAI Request/Response format. + +| Property | Details | +|----------|---------| +| Description | Vertex AI Agent Engine provides hosted agent runtimes that can execute agentic workflows with foundation models, tools, and custom logic. | +| Provider Route on LiteLLM | `vertex_ai/agent_engine/{RESOURCE_NAME}` | +| Supported Endpoints | `/chat/completions`, `/v1/messages`, `/v1/responses`, `/v1/a2a/message/send` | +| Provider Doc | [Vertex AI Agent Engine ↗](https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/overview) | + +## Quick Start + +### Model Format + +```shell showLineNumbers title="Model Format" +vertex_ai/agent_engine/{RESOURCE_NAME} +``` + +**Example:** +- `vertex_ai/agent_engine/projects/1060139831167/locations/us-central1/reasoningEngines/8263861224643493888` + +### LiteLLM Python SDK + +```python showLineNumbers title="Basic Agent Completion" +import litellm + +response = litellm.completion( + model="vertex_ai/agent_engine/projects/1060139831167/locations/us-central1/reasoningEngines/8263861224643493888", + messages=[ + {"role": "user", "content": "Explain machine learning in simple terms"} + ], +) + +print(response.choices[0].message.content) +``` + +```python showLineNumbers title="Streaming Agent Responses" +import litellm + +response = await litellm.acompletion( + model="vertex_ai/agent_engine/projects/1060139831167/locations/us-central1/reasoningEngines/8263861224643493888", + messages=[ + {"role": "user", "content": "What are the key principles of software architecture?"} + ], + stream=True, +) + +async for chunk in response: + if chunk.choices[0].delta.content: + print(chunk.choices[0].delta.content, end="") +``` + +### LiteLLM Proxy + +#### 1. Configure your model in config.yaml + + + + +```yaml showLineNumbers title="LiteLLM Proxy Configuration" +model_list: + - model_name: vertex-agent-1 + litellm_params: + model: vertex_ai/agent_engine/projects/1060139831167/locations/us-central1/reasoningEngines/8263861224643493888 + vertex_project: your-project-id + vertex_location: us-central1 +``` + + + + +#### 2. Start the LiteLLM Proxy + +```bash showLineNumbers title="Start LiteLLM Proxy" +litellm --config config.yaml +``` + +#### 3. Make requests to your Vertex AI Agent Engine + + + + +```bash showLineNumbers title="Basic Agent Request" +curl http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -d '{ + "model": "vertex-agent-1", + "messages": [ + {"role": "user", "content": "Summarize the main benefits of cloud computing"} + ] + }' +``` + + + + + +```python showLineNumbers title="Using OpenAI SDK with LiteLLM Proxy" +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:4000", + api_key="your-litellm-api-key" +) + +response = client.chat.completions.create( + model="vertex-agent-1", + messages=[ + {"role": "user", "content": "What are best practices for API design?"} + ] +) + +print(response.choices[0].message.content) +``` + + + + +## LiteLLM A2A Gateway + +You can also connect to Vertex AI Agent Engine through LiteLLM's A2A (Agent-to-Agent) Gateway UI. This provides a visual way to register and test agents without writing code. + +### 1. Navigate to Agents + +From the sidebar, click "Agents" to open the agent management page, then click "+ Add New Agent". + +![Click Agents](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/9a979927-ce6b-4168-9fba-e53e28f1c2c4/ascreenshot.jpeg?tl_px=0,14&br_px=1376,783&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=17,277) + +![Add New Agent](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/a311750c-2e85-4589-99cb-2ce7e4021e77/ascreenshot.jpeg?tl_px=0,0&br_px=1376,769&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=195,257) + +### 2. Select Vertex AI Agent Engine Type + +Click "A2A Standard" to see available agent types, then select "Vertex AI Agent Engine". + +![Select A2A Standard](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/5b1acc4c-dc3f-4639-b4a0-e64b35c228fd/ascreenshot.jpeg?tl_px=52,0&br_px=1428,769&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=524,271) + +![Select Vertex AI Agent Engine](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/2f3bab61-3e02-4db7-84f0-82200a0f4136/ascreenshot.jpeg?tl_px=0,244&br_px=1376,1013&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=477,277) + +### 3. Configure the Agent + +Fill in the following fields: + +- **Agent Name** - A friendly name for your agent (e.g., `my-vertex-agent`) +- **Reasoning Engine Resource ID** - The full resource path from Google Cloud Console (e.g., `projects/1060139831167/locations/us-central1/reasoningEngines/8263861224643493888`) +- **Vertex Project** - Your Google Cloud project ID +- **Vertex Location** - The region where your agent is deployed (e.g., `us-central1`) + +![Enter Agent Name](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/695b84c7-9511-4337-bf19-f4505ab2b72b/ascreenshot.jpeg?tl_px=0,90&br_px=1376,859&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=480,276) + +![Enter Resource ID](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/ddce64df-b3a3-4519-ab62-f137887bcea2/ascreenshot.jpeg?tl_px=0,294&br_px=1376,1063&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=440,277) + +You can find the Resource ID in Google Cloud Console under Vertex AI > Agent Engine: + +![Copy Resource ID from Google Cloud Console](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/185d7f17-cbaa-45de-948d-49d2091805ea/ascreenshot.jpeg?tl_px=0,165&br_px=1376,934&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=493,276) + +![Enter Vertex Project](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/a64da441-3e61-4811-a1e3-9f0b12c949ff/ascreenshot.jpeg?tl_px=0,233&br_px=1376,1002&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=501,277) + +You can find the Project ID in Google Cloud Console: + +![Copy Project ID from Google Cloud Console](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/9ecad3bb-a534-42d6-9604-33906014fad6/user_cropped_screenshot.webp?tl_px=0,0&br_px=1728,1028&force_format=jpeg&q=100&width=1120.0) + +![Enter Vertex Location](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/316d1f38-4fb7-4377-86b6-c0fe7ac24383/ascreenshot.jpeg?tl_px=0,330&br_px=1376,1099&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=423,277) + +### 4. Create Agent + +Click "Create Agent" to save your configuration. + +![Create Agent](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/fb04b95d-793f-4eed-acf4-d1b3b5fa65e9/ascreenshot.jpeg?tl_px=352,347&br_px=1728,1117&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=623,498) + +### 5. Test in Playground + +Go to "Playground" in the sidebar to test your agent. + +![Go to Playground](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/9e01369b-6102-4fe3-96a7-90082cadfd6e/ascreenshot.jpeg?tl_px=0,0&br_px=1376,769&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=55,226) + +### 6. Select A2A Endpoint + +Click the endpoint dropdown and select `/v1/a2a/message/send`. + +![Select Endpoint](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/d5aeac35-531b-4cf0-af2d-88f0a71fd736/ascreenshot.jpeg?tl_px=0,146&br_px=1376,915&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=299,277) + +### 7. Select Your Agent and Send a Message + +Pick your Vertex AI Agent Engine from the dropdown and send a test message. + +![Select Agent](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/353431f3-a0ba-4436-865d-ae11595e9cc4/ascreenshot.jpeg?tl_px=0,263&br_px=1376,1032&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=270,277) + +![Send Message](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/fbfce72e-f50b-43e1-b6e5-0d41192d8e2d/ascreenshot.jpeg?tl_px=95,347&br_px=1471,1117&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=524,474) + +![Agent Response](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-16/892dd826-fbf9-4530-8d82-95270889274a/ascreenshot.jpeg?tl_px=0,82&br_px=1376,851&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=485,277) + +## Environment Variables + +| Variable | Description | +|----------|-------------| +| `GOOGLE_APPLICATION_CREDENTIALS` | Path to service account JSON key file | +| `VERTEXAI_PROJECT` | Google Cloud project ID | +| `VERTEXAI_LOCATION` | Google Cloud region (default: `us-central1`) | + +```bash +export GOOGLE_APPLICATION_CREDENTIALS="/path/to/service-account.json" +export VERTEXAI_PROJECT="your-project-id" +export VERTEXAI_LOCATION="us-central1" +``` + +## Further Reading + +- [Vertex AI Agent Engine Documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/overview) +- [Create a Reasoning Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/create) +- [A2A Agent Gateway](../a2a.md) +- [Vertex AI Provider](./vertex.md) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 3c8b23f856d..c64ac1e30fe 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -632,6 +632,7 @@ const sidebars = { "providers/vertex_speech", "providers/vertex_batch", "providers/vertex_ocr", + "providers/vertex_ai_agent_engine", ] }, { diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index c5da3052738..d63b26d55fe 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1947,6 +1947,23 @@ "a2a": true } }, + "vertex_ai/agent_engine": { + "display_name": "Vertex AI Agent Engine (`vertex_ai/agent_engine`)", + "url": "https://docs.litellm.ai/docs/providers/vertex_ai_agent_engine", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true + } + }, "pydantic_ai_agents": { "display_name": "Pydantic AI Agents (`pydantic_ai_agents`)", "url": "https://docs.litellm.ai/docs/providers/pydantic_ai_agent", From 0eba879179375961a205f32ee74a1d60dd767ef7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 15 Dec 2025 20:12:44 -0800 Subject: [PATCH 040/121] Add ability to add embedding model for milvus --- .../VectorStoreForm.tsx | 94 +++++++++++++++---- .../src/components/vector_store_providers.tsx | 10 +- 2 files changed, 85 insertions(+), 19 deletions(-) diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx index d1cd3c5e443..8be879b2239 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx @@ -1,4 +1,4 @@ -import React, { useState } from "react"; +import React, { useState, useEffect } from "react"; import { TextInput, Button as TremorButton } from "@tremor/react"; import { Modal, Form, Select, Tooltip, Input, Alert } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; @@ -10,6 +10,7 @@ import { getProviderSpecificFields, VectorStoreFieldConfig, } from "../vector_store_providers"; +import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; import NotificationsManager from "../molecules/notifications_manager"; interface VectorStoreFormProps { @@ -30,6 +31,24 @@ const VectorStoreForm: React.FC = ({ const [form] = Form.useForm(); const [metadataJson, setMetadataJson] = useState("{}"); const [selectedProvider, setSelectedProvider] = useState("bedrock"); + const [modelInfo, setModelInfo] = useState([]); + + useEffect(() => { + if (!accessToken) return; + + const loadModels = async () => { + try { + const uniqueModels = await fetchAvailableModels(accessToken); + if (uniqueModels.length > 0) { + setModelInfo(uniqueModels); + } + } catch (error) { + console.error("Error fetching model info:", error); + } + }; + + loadModels(); + }, [accessToken]); const handleCreate = async (formValues: any) => { if (!accessToken) return; @@ -207,23 +226,62 @@ const VectorStoreForm: React.FC = ({ {/* Provider-specific fields */} - {getProviderSpecificFields(selectedProvider).map((field: VectorStoreFieldConfig) => ( - - {field.label}{" "} - - - - - } - name={field.name} - rules={field.required ? [{ required: true, message: `Please input the ${field.label.toLowerCase()}` }] : []} - > - - - ))} + {getProviderSpecificFields(selectedProvider).map((field: VectorStoreFieldConfig) => { + if (field.type === "select") { + const embeddingModels = modelInfo + .filter((option: ModelGroup) => option.mode === "embedding") + .map((option: ModelGroup) => ({ + value: option.model_group, + label: option.model_group, + })); + + return ( + + {field.label}{" "} + + + + + } + name={field.name} + rules={ + field.required ? [{ required: true, message: `Please select the ${field.label.toLowerCase()}` }] : [] + } + > +