From 1c8bf19f1efe35b807768a1373f8e413761f6dd8 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 20 Jan 2026 23:25:09 +0530 Subject: [PATCH] fix(proxy_server): pass search_tools to Router during DB-triggered initialization (#19388) --- litellm/proxy/proxy_server.py | 55 ++++++++++++++++++----------------- 1 file changed, 28 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c3a4de314e5..bbf434aead1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -548,9 +548,9 @@ except ImportError: server_root_path = os.getenv("SERVER_ROOT_PATH", "") _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional[ - "EnterpriseLicenseData" -] = _license_check.airgapped_license_data +premium_user_data: Optional["EnterpriseLicenseData"] = ( + _license_check.airgapped_license_data +) global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -899,7 +899,7 @@ def get_openapi_schema(): from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema) - + # Fix Swagger UI execute path error when server_root_path is set if server_root_path: openapi_schema["servers"] = [{"url": "/" + server_root_path.strip("/")}] @@ -925,7 +925,7 @@ def custom_openapi(): from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema) - + # Fix Swagger UI execute path error when server_root_path is set if server_root_path: openapi_schema["servers"] = [{"url": "/" + server_root_path.strip("/")}] @@ -1203,9 +1203,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional[ - "ClientSession" -] = None # Global shared session for connection reuse +shared_aiohttp_session: Optional["ClientSession"] = ( + None # Global shared session for connection reuse +) user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1213,9 +1213,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[ - RedisCache -] = None # redis cache used for tracking spend, tpm/rpm limits +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None @@ -1554,9 +1554,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[ - LiteLLM_TeamTable - ] = await user_api_key_cache.async_get_cache(key=_id) + existing_spend_obj: Optional[LiteLLM_TeamTable] = ( + await user_api_key_cache.async_get_cache(key=_id) + ) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -3095,11 +3095,12 @@ class ProxyConfig: async def _update_llm_router( self, - new_models: list, + new_models: Optional[Json], proxy_logging_obj: ProxyLogging, ): global llm_router, llm_model_list, master_key, general_settings - + config_data = await proxy_config.get_config() + search_tools = self.parse_search_tools(config_data) try: if llm_router is None and master_key is not None: verbose_proxy_logger.debug(f"len new_models: {len(new_models)}") @@ -3114,6 +3115,7 @@ class ProxyConfig: router_general_settings=RouterGeneralSettings( async_only_mode=True # only init async clients ), + search_tools=search_tools, ignore_invalid_deployments=True, ) verbose_proxy_logger.debug(f"updated llm_router: {llm_router}") @@ -3134,7 +3136,6 @@ class ProxyConfig: llm_model_list = llm_router.get_model_list() # check if user set any callbacks in Config Table - config_data = await proxy_config.get_config() self._add_callbacks_from_db_config(config_data) # router settings @@ -3944,10 +3945,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client + guardrails_in_db: List[Guardrail] = ( + await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -4274,9 +4275,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -9729,9 +9730,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[ - idx - ].field_description = sub_field_info.description + nested_fields[idx].field_description = ( + sub_field_info.description + ) idx += 1 _stored_in_db = None