From 1e18b33413f328ff2bd5abcca49a90e014d5dcd3 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 27 Jan 2026 16:30:43 -0800 Subject: [PATCH] perf: optimize initialize_standard_built_in_tools_params with early exit and single-pass extraction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add early exit when kwargs lacks web_search_options and tools keys - Add get_built_in_tools_from_kwargs() for single-pass tool extraction - Avoids calling _get_web_search_options() and _get_file_search_tool_call() separately (which each iterated tools list) Performance improvement: - No-tools workloads: 83.9% faster (117.9ms → 19.0ms per 6000 calls) - With-tools workloads: 14.7% faster (240.9ms → 205.5ms per 6000 calls) --- litellm/litellm_core_utils/litellm_logging.py | 14 ++- .../llm_cost_calc/tool_call_cost_tracking.py | 36 ++++++ .../test_tool_call_cost_tracking.py | 119 ++++++++++++++++++ 3 files changed, 163 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 66b14946c32..df7e9273667 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -495,13 +495,15 @@ class Logging(LiteLLMLoggingBaseClass): checks if web_search_options in kwargs or tools and sets the corresponding attribute in StandardBuiltInToolsParams """ + if not kwargs or ("web_search_options" not in kwargs and "tools" not in kwargs): + return StandardBuiltInToolsParams() + + web_search, file_search = ( + StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + ) return StandardBuiltInToolsParams( - web_search_options=StandardBuiltInToolCostTracking._get_web_search_options( - kwargs or {} - ), - file_search=StandardBuiltInToolCostTracking._get_file_search_tool_call( - kwargs or {} - ), + web_search_options=web_search, + file_search=file_search, ) def update_environment_variables( diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 4a4a2508d2e..99f269758fc 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -716,3 +716,39 @@ class StandardBuiltInToolCostTracking: if tool.get("type", None) == "file_search": return True return False + + @staticmethod + def get_built_in_tools_from_kwargs( + kwargs: Dict, + ) -> Tuple[Optional[WebSearchOptions], Optional[FileSearchTool]]: + """ + Extract web_search_options and file_search from kwargs in a single pass. + """ + web_search_options: Optional[WebSearchOptions] = None + file_search: Optional[FileSearchTool] = None + + # Check direct web_search_options first + web_search_options_dict = kwargs.get("web_search_options") + if web_search_options_dict: + web_search_options = WebSearchOptions(**web_search_options_dict) + + # Get tools once + tools = kwargs.get("tools") + if not tools: + return web_search_options, file_search + + # Single iteration through tools + for tool in tools: + if not isinstance(tool, dict): + continue + + if web_search_options is None and StandardBuiltInToolCostTracking._is_web_search_tool_call(tool): + web_search_options = WebSearchOptions(**tool) + + if file_search is None and StandardBuiltInToolCostTracking._is_file_search_tool_call(tool): + file_search = FileSearchTool(**tool) + + if web_search_options is not None and file_search is not None: + break + + return web_search_options, file_search diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 776c41697fe..de77cea92a5 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -257,3 +257,122 @@ def test_azure_assistant_features_integrated_cost_tracking(): # Note: File search integration test removed due to complex annotation detection logic # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage + + +class TestGetBuiltInToolsFromKwargs: + """Tests for get_built_in_tools_from_kwargs to ensure backwards compatibility.""" + + def test_empty_kwargs_returns_none(self): + """Empty dict should return (None, None).""" + result = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs({}) + assert result == (None, None) + + def test_kwargs_without_tools_or_web_search_returns_none(self): + """kwargs with unrelated keys should return (None, None).""" + kwargs = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} + result = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + assert result == (None, None) + + def test_direct_web_search_options_extracted(self): + """Direct web_search_options dict in kwargs should be extracted.""" + kwargs = {"web_search_options": {"search_context_size": "medium"}} + web_search, file_search = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + + assert web_search is not None + assert isinstance(web_search, dict) + assert web_search.get("search_context_size") == "medium" + assert file_search is None + + def test_web_search_tool_in_tools_list(self): + """web_search tool in tools list should be extracted.""" + kwargs = { + "tools": [ + {"type": "function", "function": {"name": "get_weather"}}, + {"type": "web_search", "search_context_size": "low"}, + ] + } + web_search, file_search = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + + assert web_search is not None + assert web_search.get("type") == "web_search" + assert file_search is None + + def test_file_search_tool_in_tools_list(self): + """file_search tool in tools list should be extracted.""" + kwargs = { + "tools": [ + {"type": "function", "function": {"name": "get_weather"}}, + {"type": "file_search", "max_results": 10}, + ] + } + web_search, file_search = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + + assert web_search is None + assert file_search is not None + assert file_search.get("type") == "file_search" + + def test_both_web_search_and_file_search_extracted(self): + """Both tool types should be extracted when present.""" + kwargs = { + "tools": [ + {"type": "web_search"}, + {"type": "file_search"}, + {"type": "function", "function": {"name": "calculate"}}, + ] + } + web_search, file_search = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + + assert web_search is not None + assert file_search is not None + + def test_only_function_tools_returns_none(self): + """Tools list with only function tools should return (None, None).""" + kwargs = { + "tools": [ + {"type": "function", "function": {"name": "get_weather"}}, + {"type": "function", "function": {"name": "send_email"}}, + ] + } + result = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + assert result == (None, None) + + def test_non_dict_tools_skipped(self): + """Non-dict items in tools list should be skipped without error.""" + kwargs = { + "tools": [ + "not a dict", + None, + 123, + {"type": "web_search"}, + ] + } + web_search, file_search = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + + assert web_search is not None + assert web_search.get("type") == "web_search" + + def test_web_search_preview_type_extracted(self): + """web_search_preview type should also be extracted as web_search.""" + kwargs = { + "tools": [ + {"type": "web_search_preview"}, + ] + } + web_search, file_search = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + + assert web_search is not None + assert web_search.get("type") == "web_search_preview" + + def test_direct_web_search_options_takes_precedence(self): + """Direct web_search_options should be used even if tools list has web_search.""" + kwargs = { + "web_search_options": {"search_context_size": "high"}, + "tools": [ + {"type": "web_search", "search_context_size": "low"}, + ] + } + web_search, file_search = StandardBuiltInToolCostTracking.get_built_in_tools_from_kwargs(kwargs) + + # Direct web_search_options should take precedence + assert web_search is not None + assert web_search.get("search_context_size") == "high"