diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md index b1166a7809c..f3a4744a062 100644 --- a/docs/my-website/docs/a2a.md +++ b/docs/my-website/docs/a2a.md @@ -257,6 +257,21 @@ LiteLLM follows the [A2A JSON-RPC 2.0 specification](https://github.com/google/A } ``` +## Agent Health Checks + +LiteLLM can automatically filter out agents whose backends are unreachable. The `GET /v1/agents` endpoint only returns agents that are currently healthy. + +- **Without background checks (default):** Health checks run inline every time `GET /v1/agents` is called. +- **With background checks:** A background loop periodically pings agents so the endpoint responds instantly. + +```yaml +general_settings: + background_agent_health_checks: true # run health checks in the background + health_check_interval: 300 # interval in seconds (default: 300) +``` + +For more details, see [Background Agent Health Checks](./proxy/health#background-agent-health-checks). + ## Agent Registry Want to create a central registry so your team can discover what agents are available within your company? diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index af868bc9f9d..1369c83a20d 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -132,6 +132,7 @@ general_settings: global_max_parallel_requests: 0 # the max parallel requests allowed on the proxy all up infer_model_from_keys: true background_health_checks: true + background_agent_health_checks: true # enable background health checks for A2A agents health_check_interval: 300 alerting: ["slack", "email"] alerting_threshold: 0 @@ -228,6 +229,7 @@ router_settings: | global_max_parallel_requests | integer | The max parallel requests allowed on the proxy overall | | infer_model_from_keys | boolean | If true, infers the model from the provided keys | | background_health_checks | boolean | If true, enables background health checks. [Doc on health checks](health) | +| background_agent_health_checks | boolean | If true, enables background health checks for A2A agents. When disabled, health checks run inline on `GET /v1/agents`. [Doc on health checks](health#background-agent-health-checks) | | health_check_interval | integer | The interval for health checks in seconds [Doc on health checks](health) | | alerting | array of strings | List of alerting methods [Doc on Slack Alerting](alerting) | | alerting_threshold | integer | The threshold for triggering alerts [Doc on Slack Alerting](alerting) | diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 2764a6f0d4f..9abf2a68d55 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -284,6 +284,30 @@ $ litellm /path/to/config.yaml curl --location 'http://0.0.0.0:4000/health' ``` +### Background Agent Health Checks + +You can enable background health checks for A2A agents registered on the proxy. When enabled, a background loop periodically pings each agent's URL and tracks which agents are healthy. The `GET /v1/agents` endpoint then only returns agents that are currently reachable. + +If `background_agent_health_checks` is **not** enabled, the health check runs inline every time `GET /v1/agents` is called. + +**How it works:** +- Sends an HTTP `HEAD` request to each agent's `url` (from `agent_card_params`) +- Any response with status code < 500 (including `405 Method Not Allowed`) means the agent's server is live +- Connection errors or timeouts mark the agent as unhealthy +- Agents using the completion bridge (no URL, only `custom_llm_provider`) are always treated as healthy + +```yaml +general_settings: + background_agent_health_checks: true # enable background agent health checks + health_check_interval: 300 # shared interval for all background health checks (seconds) +``` + +```bash +# Only healthy agents are returned +curl -X GET "http://0.0.0.0:4000/v1/agents" \ + -H "Authorization: Bearer sk-1234" +``` + ### Disable Background Health Checks For Specific Models Use this if you want to disable background health checks for specific models. diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 159c9fb93d9..a12503f4407 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -1,13 +1,12 @@ import hashlib import json from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Set import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy.management_helpers.object_permission_utils import ( - handle_update_object_permission_common, -) +from litellm.proxy.management_helpers.object_permission_utils import \ + handle_update_object_permission_common from litellm.proxy.utils import PrismaClient from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -15,6 +14,7 @@ from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest class AgentRegistry: def __init__(self): self.agent_list: List[AgentResponse] = [] + self.healthy_agent_ids: Optional[Set[str]] = None def reset_agent_list(self): self.agent_list = [] @@ -43,6 +43,43 @@ class AgentRegistry: public_agent_list.append(agent) return public_agent_list + async def run_health_check(self) -> None: + """Ping each agent's URL with a HEAD request and update healthy_agent_ids.""" + from litellm.llms.custom_httpx.http_handler import \ + get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.A2A) + healthy: Set[str] = set() + for agent in self.agent_list: + url = agent.agent_card_params.get("url") + if not url: + # Completion-bridge agents (no URL) are always treated as healthy + if agent.litellm_params and agent.litellm_params.get( + "custom_llm_provider" + ): + healthy.add(agent.agent_id) + continue + try: + resp = await client.client.head(url, timeout=10.0) + if resp.status_code < 500: # handle 405 Method Not Allowed + healthy.add(agent.agent_id) + except Exception: + pass # skip if unhealthy + self.healthy_agent_ids = healthy + + def get_healthy_agent_list( + self, agent_names: Optional[List[str]] = None + ) -> List[AgentResponse]: + """Return only agents whose health check passed. + + If no health data exists yet (healthy_agent_ids is None), returns all agents. + """ + agents = self.get_agent_list(agent_names=agent_names) + if self.healthy_agent_ids is None: + return agents + return [a for a in agents if a.agent_id in self.healthy_agent_ids] + def _create_agent_id(self, agent_config: AgentConfig) -> str: return hashlib.sha256( json.dumps(agent_config, sort_keys=True).encode() @@ -149,9 +186,13 @@ class AgentRegistry: created_agent_dict = created_agent.model_dump() if created_agent.object_permission is not None: try: - created_agent_dict["object_permission"] = created_agent.object_permission.model_dump() + created_agent_dict["object_permission"] = ( + created_agent.object_permission.model_dump() + ) except Exception: - created_agent_dict["object_permission"] = created_agent.object_permission.dict() + created_agent_dict["object_permission"] = ( + created_agent.object_permission.dict() + ) return AgentResponse(**created_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error adding agent to DB: {str(e)}") @@ -219,12 +260,10 @@ class AgentRegistry: existing_object_permission_id = existing_agent.get( "object_permission_id" ) - object_permission_id = ( - await handle_update_object_permission_common( - agent_copy, - existing_object_permission_id, - prisma_client, - ) + object_permission_id = await handle_update_object_permission_common( + agent_copy, + existing_object_permission_id, + prisma_client, ) if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id @@ -241,9 +280,13 @@ class AgentRegistry: patched_agent_dict = patched_agent.model_dump() if patched_agent.object_permission is not None: try: - patched_agent_dict["object_permission"] = patched_agent.object_permission.model_dump() + patched_agent_dict["object_permission"] = ( + patched_agent.object_permission.model_dump() + ) except Exception: - patched_agent_dict["object_permission"] = patched_agent.object_permission.dict() + patched_agent_dict["object_permission"] = ( + patched_agent.object_permission.dict() + ) return AgentResponse(**patched_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error patching agent in DB: {str(e)}") @@ -298,12 +341,10 @@ class AgentRegistry: else None ) agent_copy = dict(agent) - object_permission_id = ( - await handle_update_object_permission_common( - agent_copy, - existing_object_permission_id, - prisma_client, - ) + object_permission_id = await handle_update_object_permission_common( + agent_copy, + existing_object_permission_id, + prisma_client, ) if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id @@ -318,9 +359,13 @@ class AgentRegistry: updated_agent_dict = updated_agent.model_dump() if updated_agent.object_permission is not None: try: - updated_agent_dict["object_permission"] = updated_agent.object_permission.model_dump() + updated_agent_dict["object_permission"] = ( + updated_agent.object_permission.model_dump() + ) except Exception: - updated_agent_dict["object_permission"] = updated_agent.object_permission.dict() + updated_agent_dict["object_permission"] = ( + updated_agent.object_permission.dict() + ) return AgentResponse(**updated_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error updating agent in DB: {str(e)}") @@ -344,7 +389,9 @@ class AgentRegistry: # object_permission is eagerly loaded via include above if agent.object_permission is not None: try: - agent_dict["object_permission"] = agent.object_permission.model_dump() + agent_dict["object_permission"] = ( + agent.object_permission.model_dump() + ) except Exception: agent_dict["object_permission"] = agent.object_permission.dict() agents.append(agent_dict) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 65674d01be7..00050c1e161 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -14,19 +14,16 @@ from fastapi import APIRouter, Depends, HTTPException, Request import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import (CommonProxyErrors, LitellmUserRoles, + UserAPIKeyAuth) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity -from litellm.types.agents import ( - AgentConfig, - AgentMakePublicResponse, - AgentResponse, - MakeAgentsPublicRequest, - PatchAgentRequest, -) -from litellm.types.proxy.management_endpoints.common_daily_activity import ( - SpendAnalyticsPaginatedResponse, -) +from litellm.proxy.management_endpoints.common_daily_activity import \ + get_daily_activity +from litellm.types.agents import (AgentConfig, AgentMakePublicResponse, + AgentResponse, MakeAgentsPublicRequest, + PatchAgentRequest) +from litellm.types.proxy.management_endpoints.common_daily_activity import \ + SpendAnalyticsPaginatedResponse router = APIRouter() @@ -69,12 +66,18 @@ async def get_agents( Returns: List[AgentResponse] """ - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( - AgentRequestHandler, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ + AgentRequestHandler try: + from litellm.proxy.proxy_server import \ + use_background_agent_health_checks + + if not use_background_agent_health_checks: + await global_agent_registry.run_health_check() + returned_agents: List[AgentResponse] = [] # Admin users get all agents @@ -82,7 +85,7 @@ async def get_agents( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ): - returned_agents = global_agent_registry.get_agent_list() + returned_agents = global_agent_registry.get_healthy_agent_list() else: # Get allowed agents from object_permission (key/team level) allowed_agent_ids = await AgentRequestHandler.get_allowed_agents( @@ -91,10 +94,10 @@ async def get_agents( # If no restrictions (empty list), return all agents if len(allowed_agent_ids) == 0: - returned_agents = global_agent_registry.get_agent_list() + returned_agents = global_agent_registry.get_healthy_agent_list() else: # Filter agents by allowed IDs - all_agents = global_agent_registry.get_agent_list() + all_agents = global_agent_registry.get_healthy_agent_list() returned_agents = [ agent for agent in all_agents if agent.agent_id in allowed_agent_ids @@ -125,9 +128,8 @@ async def get_agents( #### CRUD ENDPOINTS FOR AGENTS #### -from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, -) +from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY @router.post( @@ -564,9 +566,8 @@ async def make_agent_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY from litellm.proxy.proxy_server import proxy_config # Check if user has admin permissions @@ -681,9 +682,8 @@ async def make_agents_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY from litellm.proxy.proxy_server import proxy_config # Load existing config diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 33d84cd7078..4ade5b333d2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14,22 +14,9 @@ import time import traceback import warnings from datetime import datetime, timedelta, timezone -from typing import ( - TYPE_CHECKING, - Any, - AsyncGenerator, - Dict, - List, - Literal, - Optional, - Set, - Tuple, - Union, - cast, - get_args, - get_origin, - get_type_hints, -) +from typing import (TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Literal, + Optional, Set, Tuple, Union, cast, get_args, get_origin, + get_type_hints) import anyio import websockets @@ -38,74 +25,39 @@ from pydantic import BaseModel, Json from litellm._uuid import uuid from litellm.constants import ( - AIOHTTP_CONNECTOR_LIMIT, - AIOHTTP_CONNECTOR_LIMIT_PER_HOST, - AIOHTTP_KEEPALIVE_TIMEOUT, - AIOHTTP_NEEDS_CLEANUP_CLOSED, - AIOHTTP_TTL_DNS_CACHE, - AUDIO_SPEECH_CHUNK_SIZE, - BASE_MCP_ROUTE, - DEFAULT_MAX_RECURSE_DEPTH, - DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, - DEFAULT_SHARED_HEALTH_CHECK_TTL, - DEFAULT_SLACK_ALERTING_THRESHOLD, + AIOHTTP_CONNECTOR_LIMIT, AIOHTTP_CONNECTOR_LIMIT_PER_HOST, + AIOHTTP_KEEPALIVE_TIMEOUT, AIOHTTP_NEEDS_CLEANUP_CLOSED, + AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, BASE_MCP_ROUTE, + DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, + DEFAULT_SHARED_HEALTH_CHECK_TTL, DEFAULT_SLACK_ALERTING_THRESHOLD, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, - LITELLM_SETTINGS_SAFE_DB_OVERRIDES, - LITELLM_UI_ALLOW_HEADERS, - LITELLM_UI_SESSION_DURATION, -) -from litellm.litellm_core_utils.litellm_logging import ( - _init_custom_logger_compatible_class, -) + LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, + LITELLM_UI_SESSION_DURATION) +from litellm.litellm_core_utils.litellm_logging import \ + _init_custom_logger_compatible_class from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import ( - CallbackDelete, - CallInfo, - CommonProxyErrors, - ConfigFieldDelete, - ConfigFieldInfo, - ConfigFieldUpdate, - ConfigGeneralSettings, - ConfigList, - ConfigYAML, - EnterpriseLicenseData, - FieldDetail, - InvitationClaim, - InvitationDelete, - InvitationModel, - InvitationNew, - InvitationUpdate, - Litellm_EntityType, - LiteLLM_JWTAuth, - LiteLLM_TeamTable, - LiteLLM_UserTable, - LitellmUserRoles, - PassThroughGenericEndpoint, - ProxyErrorTypes, - ProxyException, - RoleBasedPermissions, - SpecialModelNames, - SupportedDBObjectType, - TeamDefaultSettings, - TokenCountRequest, - TransformRequestBody, - UserAPIKeyAuth, -) +from litellm.proxy._types import (CallbackDelete, CallInfo, CommonProxyErrors, + ConfigFieldDelete, ConfigFieldInfo, + ConfigFieldUpdate, ConfigGeneralSettings, + ConfigList, ConfigYAML, + EnterpriseLicenseData, FieldDetail, + InvitationClaim, InvitationDelete, + InvitationModel, InvitationNew, + InvitationUpdate, Litellm_EntityType, + LiteLLM_JWTAuth, LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, + PassThroughGenericEndpoint, ProxyErrorTypes, + ProxyException, RoleBasedPermissions, + SpecialModelNames, SupportedDBObjectType, + TeamDefaultSettings, TokenCountRequest, + TransformRequestBody, UserAPIKeyAuth) from litellm.proxy.common_utils.callback_utils import ( - normalize_callback_names, - process_callback, -) + normalize_callback_names, process_callback) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body -from litellm.types.utils import ( - ModelResponse, - ModelResponseStream, - TextCompletionResponse, - TokenCountResponse, -) -from litellm.utils import ( - _invalidate_model_cost_lowercase_map, - load_credentials_from_list, -) +from litellm.types.utils import (ModelResponse, ModelResponseStream, + TextCompletionResponse, TokenCountResponse) +from litellm.utils import (_invalidate_model_cost_lowercase_map, + load_credentials_from_list) if TYPE_CHECKING: from aiohttp import ClientSession @@ -202,359 +154,272 @@ from litellm import Router from litellm._logging import verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.constants import ( - _REALTIME_BODY_CACHE_SIZE, - APSCHEDULER_COALESCE, - APSCHEDULER_MAX_INSTANCES, - APSCHEDULER_MISFIRE_GRACE_TIME, - APSCHEDULER_REPLACE_EXISTING, - DAYS_IN_A_MONTH, - DEFAULT_HEALTH_CHECK_INTERVAL, - DEFAULT_MODEL_CREATED_AT_TIME, - LITELLM_PROXY_ADMIN_NAME, - PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, - PROXY_BATCH_POLLING_INTERVAL, - PROXY_BATCH_WRITE_AT, - PROXY_BUDGET_RESCHEDULER_MAX_TIME, - PROXY_BUDGET_RESCHEDULER_MIN_TIME, -) +from litellm.constants import (_REALTIME_BODY_CACHE_SIZE, APSCHEDULER_COALESCE, + APSCHEDULER_MAX_INSTANCES, + APSCHEDULER_MISFIRE_GRACE_TIME, + APSCHEDULER_REPLACE_EXISTING, DAYS_IN_A_MONTH, + DEFAULT_HEALTH_CHECK_INTERVAL, + DEFAULT_MODEL_CREATED_AT_TIME, + LITELLM_PROXY_ADMIN_NAME, + PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, + PROXY_BATCH_POLLING_INTERVAL, + PROXY_BATCH_WRITE_AT, + PROXY_BUDGET_RESCHEDULER_MAX_TIME, + PROXY_BUDGET_RESCHEDULER_MIN_TIME) from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - get_litellm_metadata_from_kwargs, -) + _get_parent_otel_span_from_kwargs, get_litellm_metadata_from_kwargs) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.litellm_core_utils.litellm_logging import \ + Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.sensitive_data_masker import \ + SensitiveDataMasker +from litellm.llms.custom_httpx.http_handler import (AsyncHTTPHandler, + HTTPHandler) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( - router as mcp_byok_oauth_router, -) -from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - router as mcp_discoverable_endpoints_router, -) -from litellm.proxy._experimental.mcp_server.rest_endpoints import ( - router as mcp_rest_endpoints_router, -) +from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import \ + router as mcp_byok_oauth_router +from litellm.proxy._experimental.mcp_server.discoverable_endpoints import \ + router as mcp_discoverable_endpoints_router +from litellm.proxy._experimental.mcp_server.rest_endpoints import \ + router as mcp_rest_endpoints_router from litellm.proxy._experimental.mcp_server.server import app as mcp_app -from litellm.proxy._experimental.mcp_server.tool_registry import ( - global_mcp_tool_registry, -) +from litellm.proxy._experimental.mcp_server.tool_registry import \ + global_mcp_tool_registry from litellm.proxy._types import * from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry -from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router +from litellm.proxy.agent_endpoints.endpoints import \ + router as agent_endpoints_router from litellm.proxy.agent_endpoints.model_list_helpers import ( - append_agents_to_model_group, - append_agents_to_model_info, -) -from litellm.proxy.analytics_endpoints.analytics_endpoints import ( - router as analytics_router, -) -from litellm.proxy.anthropic_endpoints.claude_code_endpoints import ( - claude_code_marketplace_router, -) -from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router -from litellm.proxy.anthropic_endpoints.skills_endpoints import ( - router as anthropic_skills_router, -) -from litellm.proxy.auth.auth_checks import ( - ExperimentalUIJWTToken, - get_team_object, - log_db_metrics, -) + append_agents_to_model_group, append_agents_to_model_info) +from litellm.proxy.analytics_endpoints.analytics_endpoints import \ + router as analytics_router +from litellm.proxy.anthropic_endpoints.claude_code_endpoints import \ + claude_code_marketplace_router +from litellm.proxy.anthropic_endpoints.endpoints import \ + router as anthropic_router +from litellm.proxy.anthropic_endpoints.skills_endpoints import \ + router as anthropic_skills_router +from litellm.proxy.auth.auth_checks import (ExperimentalUIJWTToken, + get_team_object, log_db_metrics) from litellm.proxy.auth.auth_utils import check_response_size_is_safe from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import LicenseCheck -from litellm.proxy.auth.model_checks import ( - get_all_fallbacks, - get_complete_model_list, - get_key_models, - get_mcp_server_ids, - get_team_models, -) +from litellm.proxy.auth.model_checks import (get_all_fallbacks, + get_complete_model_list, + get_key_models, + get_mcp_server_ids, + get_team_models) from litellm.proxy.auth.user_api_key_auth import ( - _fetch_global_spend_with_event_coordination, - user_api_key_auth, - user_api_key_auth_websocket, -) + _fetch_global_spend_with_event_coordination, user_api_key_auth, + user_api_key_auth_websocket) from litellm.proxy.batches_endpoints.endpoints import router as batches_router - ## Import All Misc routes here ## from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - create_response, -) -from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy + ProxyBaseLLMRequestProcessing, create_response) +from litellm.proxy.common_utils.callback_utils import \ + initialize_callbacks_on_proxy from litellm.proxy.common_utils.debug_utils import init_verbose_loggers -from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router +from litellm.proxy.common_utils.debug_utils import \ + router as debugging_endpoints_router from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, - encrypt_value_helper, -) + decrypt_value_helper, encrypt_value_helper) from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.http_parsing_utils import ( - _read_request_body, - _safe_get_request_headers, - check_file_size_under_limit, - get_form_data, -) + _read_request_body, _safe_get_request_headers, check_file_size_under_limit, + get_form_data) from litellm.proxy.common_utils.load_config_utils import ( - get_config_file_contents_from_gcs, - get_file_contents_from_s3, -) -from litellm.proxy.common_utils.openai_endpoint_utils import ( - remove_sensitive_info_from_deployment, -) + get_config_file_contents_from_gcs, get_file_contents_from_s3) +from litellm.proxy.common_utils.openai_endpoint_utils import \ + remove_sensitive_info_from_deployment from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES -from litellm.proxy.container_endpoints.endpoints import router as container_router -from litellm.proxy.credential_endpoints.endpoints import router as credential_router -from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup +from litellm.proxy.container_endpoints.endpoints import \ + router as container_router +from litellm.proxy.credential_endpoints.endpoints import \ + router as credential_router +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import \ + SpendLogCleanup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router -from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router -from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config +from litellm.proxy.fine_tuning_endpoints.endpoints import \ + router as fine_tuning_router +from litellm.proxy.fine_tuning_endpoints.endpoints import \ + set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router -from litellm.proxy.guardrails.guardrail_endpoints import router as guardrails_router -from litellm.proxy.guardrails.init_guardrails import ( - init_guardrails_v2, - initialize_guardrails, -) +from litellm.proxy.guardrails.guardrail_endpoints import \ + router as guardrails_router +from litellm.proxy.guardrails.init_guardrails import (init_guardrails_v2, + initialize_guardrails) from litellm.proxy.health_check import perform_health_check -from litellm.proxy.health_endpoints._health_endpoints import router as health_router -from litellm.proxy.hooks.model_max_budget_limiter import ( - _PROXY_VirtualKeyModelMaxBudgetLimiter, -) -from litellm.proxy.hooks.prompt_injection_detection import ( - _OPTIONAL_PromptInjectionDetection, -) +from litellm.proxy.health_endpoints._health_endpoints import \ + router as health_router +from litellm.proxy.hooks.model_max_budget_limiter import \ + _PROXY_VirtualKeyModelMaxBudgetLimiter +from litellm.proxy.hooks.prompt_injection_detection import \ + _OPTIONAL_PromptInjectionDetection from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request -from litellm.proxy.management_endpoints.access_group_endpoints import ( - router as access_group_router, -) -from litellm.proxy.management_endpoints.budget_management_endpoints import ( - router as budget_management_router, -) -from litellm.proxy.management_endpoints.cache_settings_endpoints import ( - router as cache_settings_router, -) -from litellm.proxy.management_endpoints.callback_management_endpoints import ( - router as callback_management_endpoints_router, -) +from litellm.proxy.management_endpoints.access_group_endpoints import \ + router as access_group_router +from litellm.proxy.management_endpoints.budget_management_endpoints import \ + router as budget_management_router +from litellm.proxy.management_endpoints.cache_settings_endpoints import \ + router as cache_settings_router +from litellm.proxy.management_endpoints.callback_management_endpoints import \ + router as callback_management_endpoints_router from litellm.proxy.management_endpoints.common_utils import ( - _user_has_admin_privileges, - admin_can_invite_user, -) -from litellm.proxy.management_endpoints.compliance_endpoints import ( - router as compliance_router, -) -from litellm.proxy.management_endpoints.cost_tracking_settings import ( - router as cost_tracking_settings_router, -) -from litellm.proxy.management_endpoints.customer_endpoints import ( - router as customer_router, -) -from litellm.proxy.management_endpoints.fallback_management_endpoints import ( - router as fallback_management_router, -) -from litellm.proxy.management_endpoints.internal_user_endpoints import ( - router as internal_user_router, -) -from litellm.proxy.management_endpoints.internal_user_endpoints import ( - user_update, -) + _user_has_admin_privileges, admin_can_invite_user) +from litellm.proxy.management_endpoints.compliance_endpoints import \ + router as compliance_router +from litellm.proxy.management_endpoints.cost_tracking_settings import \ + router as cost_tracking_settings_router +from litellm.proxy.management_endpoints.customer_endpoints import \ + router as customer_router +from litellm.proxy.management_endpoints.fallback_management_endpoints import \ + router as fallback_management_router +from litellm.proxy.management_endpoints.internal_user_endpoints import \ + router as internal_user_router +from litellm.proxy.management_endpoints.internal_user_endpoints import \ + user_update +from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import \ + router as jwt_key_mapping_router from litellm.proxy.management_endpoints.key_management_endpoints import ( - delete_verification_tokens, - duration_in_seconds, - generate_key_helper_fn, -) -from litellm.proxy.management_endpoints.key_management_endpoints import ( - router as key_management_router, -) -from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( - router as jwt_key_mapping_router, -) -from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - router as mcp_management_router, -) -from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( - router as model_access_group_management_router, -) + delete_verification_tokens, duration_in_seconds, generate_key_helper_fn) +from litellm.proxy.management_endpoints.key_management_endpoints import \ + router as key_management_router +from litellm.proxy.management_endpoints.mcp_management_endpoints import \ + router as mcp_management_router +from litellm.proxy.management_endpoints.model_access_group_management_endpoints import \ + router as model_access_group_management_router from litellm.proxy.management_endpoints.model_management_endpoints import ( - _add_model_to_db, - _add_team_model_to_db, - _deduplicate_litellm_router_models, -) -from litellm.proxy.management_endpoints.model_management_endpoints import ( - router as model_management_router, -) -from litellm.proxy.management_endpoints.organization_endpoints import ( - router as organization_router, -) -from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router -from litellm.proxy.management_endpoints.project_endpoints import ( - router as project_router, -) -from litellm.proxy.management_endpoints.router_settings_endpoints import ( - router as router_settings_router, -) + _add_model_to_db, _add_team_model_to_db, + _deduplicate_litellm_router_models) +from litellm.proxy.management_endpoints.model_management_endpoints import \ + router as model_management_router +from litellm.proxy.management_endpoints.organization_endpoints import \ + router as organization_router +from litellm.proxy.management_endpoints.policy_endpoints import \ + router as policy_router +from litellm.proxy.management_endpoints.project_endpoints import \ + router as project_router +from litellm.proxy.management_endpoints.router_settings_endpoints import \ + router as router_settings_router from litellm.proxy.management_endpoints.scim.scim_v2 import scim_router -from litellm.proxy.management_endpoints.tag_management_endpoints import ( - router as tag_management_router, -) -from litellm.proxy.management_endpoints.team_callback_endpoints import ( - router as team_callback_router, -) -from litellm.proxy.management_endpoints.team_endpoints import router as team_router +from litellm.proxy.management_endpoints.tag_management_endpoints import \ + router as tag_management_router +from litellm.proxy.management_endpoints.team_callback_endpoints import \ + router as team_callback_router +from litellm.proxy.management_endpoints.team_endpoints import \ + router as team_router from litellm.proxy.management_endpoints.team_endpoints import ( - update_team, - validate_membership, -) -from litellm.proxy.management_endpoints.tool_management_endpoints import ( - router as tool_management_router, -) -from litellm.proxy.management_endpoints.ui_sso import ( - get_disabled_non_admin_personal_key_creation, -) + update_team, validate_membership) +from litellm.proxy.management_endpoints.tool_management_endpoints import \ + router as tool_management_router +from litellm.proxy.management_endpoints.ui_sso import \ + get_disabled_non_admin_personal_key_creation from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router -from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router -from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( - router as user_agent_analytics_router, -) -from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update -from litellm.proxy.middleware.in_flight_requests_middleware import ( - InFlightRequestsMiddleware, -) -from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware +from litellm.proxy.management_endpoints.usage_endpoints import \ + router as usage_ai_router +from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import \ + router as user_agent_analytics_router +from litellm.proxy.management_helpers.audit_logs import \ + create_audit_log_for_update +from litellm.proxy.middleware.in_flight_requests_middleware import \ + InFlightRequestsMiddleware +from litellm.proxy.middleware.prometheus_auth_middleware import \ + PrometheusAuthMiddleware from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router -from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_router -from litellm.proxy.openai_files_endpoints.files_endpoints import ( - router as openai_files_router, -) -from litellm.proxy.openai_files_endpoints.files_endpoints import ( - set_files_config, -) -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - passthrough_endpoint_router, -) -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - router as llm_passthrough_router, -) -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - vertex_ai_live_websocket_passthrough, -) -from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - initialize_pass_through_endpoints, -) -from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - router as pass_through_router, -) -from litellm.proxy.policy_engine.policy_endpoints import router as policy_crud_router -from litellm.proxy.policy_engine.policy_resolve_endpoints import ( - router as policy_resolve_router, -) +from litellm.proxy.openai_evals_endpoints.endpoints import \ + router as evals_router +from litellm.proxy.openai_files_endpoints.files_endpoints import \ + router as openai_files_router +from litellm.proxy.openai_files_endpoints.files_endpoints import \ + set_files_config +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import \ + passthrough_endpoint_router +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import \ + router as llm_passthrough_router +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import \ + vertex_ai_live_websocket_passthrough +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import \ + initialize_pass_through_endpoints +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import \ + router as pass_through_router +from litellm.proxy.policy_engine.policy_endpoints import \ + router as policy_crud_router +from litellm.proxy.policy_engine.policy_resolve_endpoints import \ + router as policy_resolve_router from litellm.proxy.prompts.prompt_endpoints import router as prompts_router from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router -from litellm.proxy.response_api_endpoints.endpoints import router as response_router +from litellm.proxy.response_api_endpoints.endpoints import \ + router as response_router from litellm.proxy.route_llm_request import route_request from litellm.proxy.search_endpoints.endpoints import router as search_router -from litellm.proxy.search_endpoints.search_tool_management import ( - router as search_tool_management_router, -) -from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router -from litellm.proxy.spend_tracking.spend_management_endpoints import ( - router as spend_management_router, -) -from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.search_endpoints.search_tool_management import \ + router as search_tool_management_router +from litellm.proxy.spend_tracking.cloudzero_endpoints import \ + router as cloudzero_router +from litellm.proxy.spend_tracking.spend_management_endpoints import \ + router as spend_management_router +from litellm.proxy.spend_tracking.spend_tracking_utils import \ + get_logging_payload from litellm.proxy.types_utils.utils import get_instance_fn -from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( - router as ui_crud_endpoints_router, -) -from litellm.proxy.utils import ( - PrismaClient, - ProxyLogging, - ProxyUpdateSpend, - _cache_user_row, - _get_docs_url, - _get_projected_spend_over_limit, - _get_redoc_url, - _is_projected_spend_over_limit, - _is_valid_team_configs, - get_custom_url, - get_error_message_str, - get_server_root_path, - handle_exception_on_proxy, - hash_token, - model_dump_with_preserved_fields, - update_spend, -) -from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router -from litellm.proxy.vector_store_endpoints.management_endpoints import ( - router as vector_store_management_router, -) -from litellm.proxy.vector_store_files_endpoints.endpoints import ( - router as vector_store_files_router, -) -from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import ( - router as langfuse_router, -) +from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import \ + router as ui_crud_endpoints_router +from litellm.proxy.utils import (PrismaClient, ProxyLogging, ProxyUpdateSpend, + _cache_user_row, _get_docs_url, + _get_projected_spend_over_limit, + _get_redoc_url, + _is_projected_spend_over_limit, + _is_valid_team_configs, get_custom_url, + get_error_message_str, get_server_root_path, + handle_exception_on_proxy, hash_token, + model_dump_with_preserved_fields, + update_spend) +from litellm.proxy.vector_store_endpoints.endpoints import \ + router as vector_store_router +from litellm.proxy.vector_store_endpoints.management_endpoints import \ + router as vector_store_management_router +from litellm.proxy.vector_store_files_endpoints.endpoints import \ + router as vector_store_files_router +from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import \ + router as langfuse_router from litellm.proxy.video_endpoints.endpoints import router as video_router -from litellm.router import ( - AssistantsTypedDict, - Deployment, - LiteLLM_Params, - ModelGroupInfo, -) +from litellm.router import (AssistantsTypedDict, Deployment, LiteLLM_Params, + ModelGroupInfo) from litellm.scheduler import FlowItem, Scheduler from litellm.secret_managers.aws_secret_manager import load_aws_kms from litellm.secret_managers.google_kms import load_google_kms -from litellm.secret_managers.main import ( - get_secret, - get_secret_bool, - get_secret_str, - str_to_bool, -) +from litellm.secret_managers.main import (get_secret, get_secret_bool, + get_secret_str, str_to_bool) from litellm.types.integrations.slack_alerting import SlackAlertingArgs -from litellm.types.llms.anthropic import ( - AnthropicMessagesRequest, - AnthropicResponse, - AnthropicResponseContentBlockText, - AnthropicResponseUsageBlock, -) +from litellm.types.llms.anthropic import (AnthropicMessagesRequest, + AnthropicResponse, + AnthropicResponseContentBlockText, + AnthropicResponseUsageBlock) from litellm.types.llms.openai import HttpxBinaryResponseContent -from litellm.types.proxy.management_endpoints.model_management_endpoints import ( - ModelGroupInfoProxy, -) +from litellm.types.proxy.management_endpoints.model_management_endpoints import \ + ModelGroupInfoProxy from litellm.types.proxy.management_endpoints.ui_sso import ( - DefaultTeamSSOParams, - LiteLLM_UpperboundKeyGenerateParams, -) + DefaultTeamSSOParams, LiteLLM_UpperboundKeyGenerateParams) from litellm.types.realtime import RealtimeQueryParams -from litellm.types.router import ( - DeploymentTypedDict, -) +from litellm.types.router import DeploymentTypedDict from litellm.types.router import ModelInfo as RouterModelInfo -from litellm.types.router import ( - RouterGeneralSettings, - SearchToolTypedDict, - updateDeployment, -) +from litellm.types.router import (RouterGeneralSettings, SearchToolTypedDict, + updateDeployment) from litellm.types.scheduler import DefaultPriorities -from litellm.types.secret_managers.main import ( - KeyManagementSettings, - KeyManagementSystem, -) +from litellm.types.secret_managers.main import (KeyManagementSettings, + KeyManagementSystem) from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.types.utils import RawRequestTypedDict, StandardLoggingPayload @@ -568,34 +433,15 @@ litellm.suppress_debug_info = True import json from typing import Union -from fastapi import ( - Depends, - FastAPI, - File, - Form, - Header, - HTTPException, - Path, - Query, - Request, - Response, - UploadFile, - WebSocket, - WebSocketDisconnect, - applications, - status, -) +from fastapi import (Depends, FastAPI, File, Form, Header, HTTPException, Path, + Query, Request, Response, UploadFile, WebSocket, + WebSocketDisconnect, applications, status) from fastapi.encoders import jsonable_encoder from fastapi.middleware.cors import CORSMiddleware from fastapi.openapi.docs import get_swagger_ui_html from fastapi.openapi.utils import get_openapi -from fastapi.responses import ( - FileResponse, - JSONResponse, - ORJSONResponse, - RedirectResponse, - StreamingResponse, -) +from fastapi.responses import (FileResponse, JSONResponse, ORJSONResponse, + RedirectResponse, StreamingResponse) from fastapi.routing import APIRouter from fastapi.security import OAuth2PasswordBearer from fastapi.security.api_key import APIKeyHeader @@ -618,7 +464,8 @@ except Exception: ################### # Import enterprise routes try: - from litellm_enterprise.proxy.enterprise_routes import router as _enterprise_router + from litellm_enterprise.proxy.enterprise_routes import \ + router as _enterprise_router from litellm_enterprise.proxy.proxy_server import EnterpriseProxyConfig enterprise_router = _enterprise_router @@ -676,7 +523,7 @@ _description = ( def cleanup_router_config_variables(): - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, health_check_concurrency, prisma_client + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, use_background_agent_health_checks, health_check_interval, health_check_concurrency, prisma_client # Set all variables to None master_key = None @@ -689,6 +536,7 @@ def cleanup_router_config_variables(): user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None use_shared_health_check = None + use_background_agent_health_checks = None health_check_interval = None health_check_concurrency = None prisma_client = None @@ -894,6 +742,11 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 _run_background_health_check() ) # start the background health check coroutine. + if use_background_agent_health_checks: + asyncio.create_task( + _run_background_agent_health_check() + ) + ## [Optional] Initialize dd tracer ProxyStartupEvent._init_dd_tracer() @@ -960,9 +813,8 @@ def get_openapi_schema(): return app.openapi_schema # Use compatibility wrapper for FastAPI 0.120+ schema generation - from litellm.proxy.common_utils.openapi_schema_compat import ( - get_openapi_schema_with_compat, - ) + from litellm.proxy.common_utils.openapi_schema_compat import \ + get_openapi_schema_with_compat openapi_schema = get_openapi_schema_with_compat( get_openapi_func=get_openapi, @@ -1016,7 +868,8 @@ def get_openapi_schema(): } # Add LLM API request schema bodies for documentation - from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec + from litellm.proxy.common_utils.custom_openapi_spec import \ + CustomOpenAPISpec openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema) @@ -1042,7 +895,8 @@ def custom_openapi(): openapi_schema["paths"] = paths_to_include # Add LLM API request schema bodies for documentation - from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec + from litellm.proxy.common_utils.custom_openapi_spec import \ + CustomOpenAPISpec openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema) @@ -1498,6 +1352,7 @@ user_custom_sso = None user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None use_shared_health_check = None +use_background_agent_health_checks = None use_queue = False health_check_interval = None health_check_concurrency = None @@ -2036,9 +1891,8 @@ def _schedule_background_health_check_db_save( return import time as time_module - from litellm.proxy.health_endpoints._health_endpoints import ( - _save_background_health_checks_to_db, - ) + from litellm.proxy.health_endpoints._health_endpoints import \ + _save_background_health_checks_to_db checked_by = ( shared_health_manager.pod_id @@ -2058,6 +1912,23 @@ def _schedule_background_health_check_db_save( ) +async def _run_background_agent_health_check(): + """Periodically run health checks on registered A2A agents.""" + global health_check_interval + + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + + while True: + try: + await global_agent_registry.run_health_check() + except Exception as e: + verbose_proxy_logger.exception( + "background agent health check failed: %s", e + ) + await asyncio.sleep(health_check_interval or 300) + + async def _run_background_health_check(): """ Periodically run health checks in the background on the endpoints. @@ -2098,9 +1969,8 @@ async def _run_background_health_check(): # Initialize shared health check manager if Redis is available and feature is enabled shared_health_manager = None if use_shared_health_check and redis_usage_cache is not None: - from litellm.proxy.health_check_utils.shared_health_check_manager import ( - SharedHealthCheckManager, - ) + from litellm.proxy.health_check_utils.shared_health_check_manager import \ + SharedHealthCheckManager shared_health_manager = SharedHealthCheckManager( redis_cache=redis_usage_cache, @@ -2678,7 +2548,7 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, health_check_concurrency, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, proxy_batch_polling_interval, config_passthrough_endpoints + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, use_background_agent_health_checks, health_check_interval, health_check_concurrency, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, proxy_batch_polling_interval, config_passthrough_endpoints config: dict = await self.get_config(config_file_path=config_file_path) @@ -2784,25 +2654,24 @@ class ProxyConfig: litellm.guardrail_name_config_map = guardrail_name_config_map elif key == "global_prompt_directory": - from litellm.integrations.dotprompt import ( - set_global_prompt_directory, - ) + from litellm.integrations.dotprompt import \ + set_global_prompt_directory set_global_prompt_directory(value) verbose_proxy_logger.info( f"{blue_color_code}Set Global Prompt Directory on LiteLLM Proxy{reset_color_code}" ) elif key == "global_bitbucket_config": - from litellm.integrations.bitbucket import ( - set_global_bitbucket_config, - ) + from litellm.integrations.bitbucket import \ + set_global_bitbucket_config set_global_bitbucket_config(value) verbose_proxy_logger.info( f"{blue_color_code}Set Global BitBucket Config on LiteLLM Proxy{reset_color_code}" ) elif key == "global_gitlab_config": - from litellm.integrations.gitlab import set_global_gitlab_config + from litellm.integrations.gitlab import \ + set_global_gitlab_config set_global_gitlab_config(value) verbose_proxy_logger.info( @@ -2873,9 +2742,8 @@ class ProxyConfig: callback ) if "prometheus" in callback: - from litellm.integrations.prometheus import ( - PrometheusLogger, - ) + from litellm.integrations.prometheus import \ + PrometheusLogger if PrometheusLogger is not None: verbose_proxy_logger.debug( @@ -3100,6 +2968,9 @@ class ProxyConfig: use_background_health_checks = general_settings.get( "background_health_checks", False ) + use_background_agent_health_checks = general_settings.get( + "background_agent_health_checks", False + ) # Enable shared health check state across pods (requires Redis) use_shared_health_check = general_settings.get( "use_shared_health_check", False @@ -3286,9 +3157,8 @@ class ProxyConfig: mcp_servers_config = config.get("mcp_servers", None) if mcp_servers_config: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ + global_mcp_server_manager # Get mcp_aliases from litellm_settings if available litellm_settings = config.get("litellm_settings", {}) @@ -3301,7 +3171,8 @@ class ProxyConfig: ## VECTOR STORES vector_store_registry_config = config.get("vector_store_registry", None) if vector_store_registry_config: - from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + from litellm.vector_stores.vector_store_registry import \ + VectorStoreRegistry if litellm.vector_store_registry is None: litellm.vector_store_registry = VectorStoreRegistry() @@ -3328,7 +3199,8 @@ class ProxyConfig: """ from litellm.proxy.policy_engine.init_policies import init_policies - from litellm.proxy.policy_engine.policy_validator import PolicyValidator + from litellm.proxy.policy_engine.policy_validator import \ + PolicyValidator if config is None: verbose_proxy_logger.debug("Policy engine: config is None, skipping") @@ -3417,9 +3289,8 @@ class ProxyConfig: key_management_system == KeyManagementSystem.AWS_SECRET_MANAGER.value # noqa: F405 ): - from litellm.secret_managers.aws_secret_manager_v2 import ( - AWSSecretsManagerV2, - ) + from litellm.secret_managers.aws_secret_manager_v2 import \ + AWSSecretsManagerV2 AWSSecretsManagerV2.load_aws_secret_manager( use_aws_secret_manager=True, @@ -3430,28 +3301,24 @@ class ProxyConfig: elif ( key_management_system == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value ): - from litellm.secret_managers.google_secret_manager import ( - GoogleSecretManager, - ) + from litellm.secret_managers.google_secret_manager import \ + GoogleSecretManager GoogleSecretManager() elif key_management_system == KeyManagementSystem.HASHICORP_VAULT.value: - from litellm.secret_managers.hashicorp_secret_manager import ( - HashicorpSecretManager, - ) + from litellm.secret_managers.hashicorp_secret_manager import \ + HashicorpSecretManager HashicorpSecretManager() elif key_management_system == KeyManagementSystem.CYBERARK.value: - from litellm.secret_managers.cyberark_secret_manager import ( - CyberArkSecretManager, - ) + from litellm.secret_managers.cyberark_secret_manager import \ + CyberArkSecretManager CyberArkSecretManager() elif key_management_system == KeyManagementSystem.CUSTOM.value: ### LOAD CUSTOM SECRET MANAGER ### - from litellm.secret_managers.custom_secret_manager_loader import ( - load_custom_secret_manager, - ) + from litellm.secret_managers.custom_secret_manager_loader import \ + load_custom_secret_manager load_custom_secret_manager(config_file_path=config_file_path) else: @@ -4033,9 +3900,8 @@ class ProxyConfig: # Schedule new job if retention period is set (not None) retention_period = general_settings.get("maximum_spend_logs_retention_period") if retention_period is not None: - from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( - SpendLogCleanup, - ) + from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import \ + SpendLogCleanup spend_log_cleanup = SpendLogCleanup() cleanup_cron = general_settings.get("maximum_spend_logs_cleanup_cron") @@ -4062,9 +3928,8 @@ class ProxyConfig: ) else: # Interval-based scheduling (existing behavior) - from litellm.litellm_core_utils.duration_parser import ( - duration_in_seconds, - ) + from litellm.litellm_core_utils.duration_parser import \ + duration_in_seconds retention_interval = general_settings.get( "maximum_spend_logs_retention_interval", "1d" @@ -4431,9 +4296,8 @@ class ProxyConfig: if self._should_load_db_object(object_type="sso_settings"): await self._init_sso_settings_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="cache_settings"): - from litellm.proxy.management_endpoints.cache_settings_endpoints import ( - CacheSettingsManager, - ) + from litellm.proxy.management_endpoints.cache_settings_endpoints import \ + CacheSettingsManager await CacheSettingsManager.init_cache_settings_in_db( prisma_client=prisma_client, proxy_config=self @@ -4450,7 +4314,8 @@ class ProxyConfig: import json import litellm - from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.proxy.hooks.mcp_semantic_filter import \ + SemanticToolFilterHook try: # Load litellm_settings from DB @@ -4615,9 +4480,8 @@ class ProxyConfig: if should_reload: # Perform the reload - from litellm.litellm_core_utils.get_model_cost_map import ( - get_model_cost_map, - ) + from litellm.litellm_core_utils.get_model_cost_map import \ + get_model_cost_map model_cost_map_url = litellm.model_cost_map_url new_model_cost_map = get_model_cost_map(url=model_cost_map_url) @@ -4723,9 +4587,8 @@ class ProxyConfig: if should_reload: # Perform the reload - from litellm.anthropic_beta_headers_manager import ( - reload_beta_headers_config, - ) + from litellm.anthropic_beta_headers_manager import \ + reload_beta_headers_config new_config = reload_beta_headers_config() @@ -4776,12 +4639,14 @@ class ProxyConfig: Returns: The PromptSpec object """ - from litellm.proxy.prompts.prompt_endpoints import create_versioned_prompt_spec + from litellm.proxy.prompts.prompt_endpoints import \ + create_versioned_prompt_spec return create_versioned_prompt_spec(db_prompt=db_prompt) async def _init_prompts_in_db(self, prisma_client: PrismaClient): - from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.prompts.prompt_registry import \ + IN_MEMORY_PROMPT_REGISTRY from litellm.types.prompts.init_prompts import PromptSpec try: @@ -4799,10 +4664,7 @@ class ProxyConfig: async def _init_guardrails_in_db(self, prisma_client: PrismaClient): from litellm.proxy.guardrails.guardrail_registry import ( - IN_MEMORY_GUARDRAIL_HANDLER, - Guardrail, - GuardrailRegistry, - ) + IN_MEMORY_GUARDRAIL_HANDLER, Guardrail, GuardrailRegistry) try: guardrails_in_db: List[ @@ -4828,10 +4690,10 @@ class ProxyConfig: """ Initialize policies and policy attachments from database into the in-memory registries. """ - from litellm.proxy.policy_engine.attachment_registry import ( - get_attachment_registry, - ) - from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.policy_engine.attachment_registry import \ + get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import \ + get_policy_registry try: # Get the global singleton instances @@ -4861,7 +4723,8 @@ class ProxyConfig: Initialize tool policy from database into the in-memory registry. Synced periodically by add_deployment -> _init_non_llm_objects_in_db. """ - from litellm.proxy.db.tool_registry_writer import get_tool_policy_registry + from litellm.proxy.db.tool_registry_writer import \ + get_tool_policy_registry try: registry = get_tool_policy_registry() @@ -4875,7 +4738,8 @@ class ProxyConfig: ) async def _init_vector_stores_in_db(self, prisma_client: PrismaClient): - from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + from litellm.vector_stores.vector_store_registry import \ + VectorStoreRegistry try: # read vector stores from db table @@ -4902,7 +4766,8 @@ class ProxyConfig: ) async def _init_vector_store_indexes_in_db(self, prisma_client: PrismaClient): - from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry + from litellm.vector_stores.vector_store_registry import \ + VectorStoreIndexRegistry try: # read vector stores from db table @@ -4932,7 +4797,8 @@ class ProxyConfig: ) async def _init_mcp_servers_in_db(self): - from litellm.proxy._experimental.mcp_server.utils import is_mcp_available + from litellm.proxy._experimental.mcp_server.utils import \ + is_mcp_available if not is_mcp_available(): verbose_proxy_logger.debug( @@ -4940,9 +4806,8 @@ class ProxyConfig: ) return - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ + global_mcp_server_manager try: await global_mcp_server_manager.reload_servers_from_database() @@ -4954,9 +4819,8 @@ class ProxyConfig: ) async def _init_agents_in_db(self, prisma_client: PrismaClient): - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY try: db_agents = await AGENT_REGISTRY.get_all_agents_from_db( @@ -4979,9 +4843,8 @@ class ProxyConfig: """ global llm_router - from litellm.proxy.search_endpoints.search_tool_registry import ( - SearchToolRegistry, - ) + from litellm.proxy.search_endpoints.search_tool_registry import \ + SearchToolRegistry from litellm.router_utils.search_api_router import SearchAPIRouter try: @@ -5021,9 +4884,8 @@ class ProxyConfig: ) async def _init_pass_through_endpoints_in_db(self): - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - initialize_pass_through_endpoints_in_db, - ) + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import \ + initialize_pass_through_endpoints_in_db await initialize_pass_through_endpoints_in_db() @@ -5121,11 +4983,8 @@ async def initialize( # noqa: PLR0915 if debug is True: # this needs to be first, so users can see Router init debugg import logging - from litellm._logging import ( - verbose_logger, - verbose_proxy_logger, - verbose_router_logger, - ) + from litellm._logging import (verbose_logger, verbose_proxy_logger, + verbose_router_logger) # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS verbose_logger.setLevel(level=logging.INFO) # sets package logs to info @@ -5134,11 +4993,8 @@ async def initialize( # noqa: PLR0915 if detailed_debug is True: import logging - from litellm._logging import ( - verbose_logger, - verbose_proxy_logger, - verbose_router_logger, - ) + from litellm._logging import (verbose_logger, verbose_proxy_logger, + verbose_router_logger) verbose_logger.setLevel(level=logging.DEBUG) # set package log to debug verbose_router_logger.setLevel(level=logging.DEBUG) # set router logs to debug @@ -5150,7 +5006,8 @@ async def initialize( # noqa: PLR0915 if litellm_log_setting.upper() == "INFO": import logging - from litellm._logging import verbose_proxy_logger, verbose_router_logger + from litellm._logging import (verbose_proxy_logger, + verbose_router_logger) # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS @@ -5163,11 +5020,9 @@ async def initialize( # noqa: PLR0915 elif litellm_log_setting.upper() == "DEBUG": import logging - from litellm._logging import ( - verbose_logger, - verbose_proxy_logger, - verbose_router_logger, - ) + from litellm._logging import (verbose_logger, + verbose_proxy_logger, + verbose_router_logger) verbose_logger.setLevel(level=logging.DEBUG) # set package log to debug verbose_router_logger.setLevel( @@ -5521,7 +5376,8 @@ class ProxyStartupEvent: litellm_settings: Dict[str, Any], ): """Initialize MCP semantic tool filter if configured""" - from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.proxy.hooks.mcp_semantic_filter import \ + SemanticToolFilterHook mcp_semantic_filter_config = litellm_settings.get( "mcp_semantic_tool_filter", None @@ -5633,9 +5489,8 @@ class ProxyStartupEvent: _teams = litellm.default_internal_user_params.get("teams") or [] if _teams and all(isinstance(team, dict) for team in _teams): - from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( - update_default_team_member_budget, - ) + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import \ + update_default_team_member_budget teams_pydantic_obj = [NewUserRequestTeam(**team) for team in _teams] await update_default_team_member_budget( @@ -5856,9 +5711,8 @@ class ProxyStartupEvent: ### CHECK BATCH COST ### if llm_router is not None: try: - from litellm_enterprise.proxy.common_utils.check_batch_cost import ( - CheckBatchCost, - ) + from litellm_enterprise.proxy.common_utils.check_batch_cost import \ + CheckBatchCost check_batch_cost_job = CheckBatchCost( proxy_logging_obj=proxy_logging_obj, @@ -5887,9 +5741,8 @@ class ProxyStartupEvent: ### CHECK RESPONSES COST ### if llm_router is not None: try: - from litellm_enterprise.proxy.common_utils.check_responses_cost import ( - CheckResponsesCost, - ) + from litellm_enterprise.proxy.common_utils.check_responses_cost import \ + CheckResponsesCost check_responses_cost_job = CheckResponsesCost( proxy_logging_obj=proxy_logging_obj, @@ -5949,7 +5802,8 @@ class ProxyStartupEvent: ######################################################## from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger from litellm.integrations.focus.focus_logger import FocusLogger - from litellm.proxy.spend_tracking.cloudzero_endpoints import is_cloudzero_setup + from litellm.proxy.spend_tracking.cloudzero_endpoints import \ + is_cloudzero_setup if await is_cloudzero_setup(): await CloudZeroLogger.init_cloudzero_background_job(scheduler=scheduler) @@ -5971,17 +5825,15 @@ class ProxyStartupEvent: ######################################################## from litellm.constants import ( LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS, - LITELLM_KEY_ROTATION_ENABLED, - ) + LITELLM_KEY_ROTATION_ENABLED) key_rotation_enabled: Optional[bool] = str_to_bool(LITELLM_KEY_ROTATION_ENABLED) verbose_proxy_logger.debug(f"key_rotation_enabled: {key_rotation_enabled}") if key_rotation_enabled is True: try: - from litellm.proxy.common_utils.key_rotation_manager import ( - KeyRotationManager, - ) + from litellm.proxy.common_utils.key_rotation_manager import \ + KeyRotationManager # Get prisma_client from global scope global prisma_client @@ -6142,9 +5994,7 @@ class ProxyStartupEvent: Doc: https://docs.datadoghq.com/tracing/trace_collection/automatic_instrumentation/dd_libraries/python/ """ from litellm.litellm_core_utils.dd_tracing import ( - _should_use_dd_profiler, - _should_use_dd_tracer, - ) + _should_use_dd_profiler, _should_use_dd_tracer) if _should_use_dd_tracer(): import ddtrace @@ -6256,13 +6106,10 @@ async def model_list( """ global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj - from litellm.proxy.management_endpoints.common_utils import ( - _user_has_admin_privileges, - ) - from litellm.proxy.utils import ( - create_model_info_response, - get_available_models_for_user, - ) + from litellm.proxy.management_endpoints.common_utils import \ + _user_has_admin_privileges + from litellm.proxy.utils import (create_model_info_response, + get_available_models_for_user) # Validate scope parameter if provided if scope is not None and scope != "expand": @@ -6389,11 +6236,9 @@ async def model_info( """ global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj - from litellm.proxy.utils import ( - create_model_info_response, - get_available_models_for_user, - validate_model_access, - ) + from litellm.proxy.utils import (create_model_info_response, + get_available_models_for_user, + validate_model_access) # Get available models for the user all_models = await get_available_models_for_user( @@ -8347,7 +8192,8 @@ def _get_provider_token_counter( if deployment is None: return None - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + from litellm.litellm_core_utils.get_llm_provider_logic import \ + get_llm_provider full_model = deployment.get("litellm_params", {}).get("model", "") model: Optional[str] = None @@ -10703,7 +10549,8 @@ async def fallback_login(request: Request): ) # hidden since this is a helper for UI sso login async def login(request: Request): # noqa: PLR0915 global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import (authenticate_user, + create_ui_token_object) from litellm.proxy.utils import get_custom_url form = await request.form() @@ -10753,7 +10600,8 @@ async def login(request: Request): # noqa: PLR0915 ) # hidden helper for UI logins via API async def login_v2(request: Request): # noqa: PLR0915 global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import (authenticate_user, + create_ui_token_object) from litellm.proxy.utils import get_custom_url try: @@ -11082,7 +10930,8 @@ async def get_image(): if logo_path.startswith(("http://", "https://")): try: # Download the image and cache it - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.llms.custom_httpx.http_handler import \ + get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider async_client = get_async_httpx_client( @@ -11126,7 +10975,8 @@ async def get_favicon(): if favicon_url.startswith(("http://", "https://")): try: - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.llms.custom_httpx.http_handler import \ + get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider async_client = get_async_httpx_client( @@ -11191,9 +11041,8 @@ async def new_invitation( ``` """ try: - from litellm.proxy.management_helpers.user_invitation import ( - create_invitation_for_user, - ) + from litellm.proxy.management_helpers.user_invitation import \ + create_invitation_for_user global prisma_client @@ -12280,7 +12129,8 @@ async def reload_model_cost_map( ) # Immediately reload the model cost map in the current pod - from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + from litellm.litellm_core_utils.get_model_cost_map import \ + get_model_cost_map model_cost_map_url = litellm.model_cost_map_url new_model_cost_map = get_model_cost_map(url=model_cost_map_url) @@ -12579,9 +12429,8 @@ async def get_model_cost_map_source( ) try: - from litellm.litellm_core_utils.get_model_cost_map import ( - get_model_cost_map_source_info, - ) + from litellm.litellm_core_utils.get_model_cost_map import \ + get_model_cost_map_source_info source_info = get_model_cost_map_source_info() model_count = len(litellm.model_cost) if litellm.model_cost else 0 @@ -12633,7 +12482,8 @@ async def reload_anthropic_beta_headers( ) # Immediately reload the beta headers config in the current pod - from litellm.anthropic_beta_headers_manager import reload_beta_headers_config + from litellm.anthropic_beta_headers_manager import \ + reload_beta_headers_config new_config = reload_beta_headers_config() @@ -13042,9 +12892,8 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): """Handle dynamic MCP server routes like /github_mcp/mcp""" try: # Validate that the MCP server exists - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ + global_mcp_server_manager from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.types.mcp import MCPAuth @@ -13063,9 +12912,8 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): scope["path"] = f"/mcp/{mcp_server_name}" # Import the MCP handler - from litellm.proxy._experimental.mcp_server.server import ( - handle_streamable_http_mcp, - ) + from litellm.proxy._experimental.mcp_server.server import \ + handle_streamable_http_mcp # Create a custom send function to capture the response response_started = False