vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix LiteLLM callback leak: cache router handlers by llm_key (#4678)

Browse Source
This commit is contained in:
pedrohsdb
2026-02-09 19:42:46 -08:00
committed by GitHub
parent 3e3848628c
commit ab292175a5
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
skyvern/forge/sdk/api/llm/api_handler_factory.py
View File

@@ -165,6 +165,7 @@ def _convert_allowed_fails_policy(policy: LLMAllowedFailsPolicy | None) -> Allow
class LLMAPIHandlerFactory: class LLMAPIHandlerFactory:
_custom_handlers: dict[str, LLMAPIHandler] = {} _custom_handlers: dict[str, LLMAPIHandler] = {}
_router_handler_cache: dict[str, LLMAPIHandler] = {}
_thinking_budget_settings: dict[str, int] | None = None _thinking_budget_settings: dict[str, int] | None = None
_prompt_caching_settings: dict[str, bool] | None = None _prompt_caching_settings: dict[str, bool] | None = None
@@ -389,6 +390,9 @@ class LLMAPIHandlerFactory:
@staticmethod @staticmethod
def get_llm_api_handler_with_router(llm_key: str) -> LLMAPIHandler: def get_llm_api_handler_with_router(llm_key: str) -> LLMAPIHandler:
if llm_key in LLMAPIHandlerFactory._router_handler_cache:
return LLMAPIHandlerFactory._router_handler_cache[llm_key]
llm_config = LLMConfigRegistry.get_config(llm_key) llm_config = LLMConfigRegistry.get_config(llm_key)
if not isinstance(llm_config, LLMRouterConfig): if not isinstance(llm_config, LLMRouterConfig):
raise InvalidLLMConfigError(llm_key) raise InvalidLLMConfigError(llm_key)
@@ -831,6 +835,7 @@ class LLMAPIHandlerFactory:
LOG.error("Failed to persist artifacts", exc_info=True) LOG.error("Failed to persist artifacts", exc_info=True)
llm_api_handler_with_router_and_fallback.llm_key = llm_key # type: ignore[attr-defined] llm_api_handler_with_router_and_fallback.llm_key = llm_key # type: ignore[attr-defined]
LLMAPIHandlerFactory._router_handler_cache[llm_key] = llm_api_handler_with_router_and_fallback
return llm_api_handler_with_router_and_fallback return llm_api_handler_with_router_and_fallback
@staticmethod @staticmethod