mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
perf: optimize initialize_standard_built_in_tools_params with early exit and single-pass extraction
- 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)
This commit is contained in:
parent
51339f5ef1
commit
1e18b33413
3 changed files with 163 additions and 6 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue