mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
working: async_post_call_response_headers_hook
This commit is contained in:
parent
6751c73c43
commit
3967326922
1 changed files with 28 additions and 18 deletions
|
|
@ -94,6 +94,11 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
|
||||
return openai_tools
|
||||
|
||||
def _get_metadata_variable_name(self, data: dict) -> str:
|
||||
if "litellm_metadata" in data:
|
||||
return "litellm_metadata"
|
||||
return "metadata"
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
|
|
@ -152,12 +157,6 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
f"Expanded {len(tools)} MCP reference(s) to {len(expanded_tools)} tools"
|
||||
)
|
||||
|
||||
# Debug: log first expanded tool format
|
||||
if expanded_tools:
|
||||
verbose_proxy_logger.debug(
|
||||
f"First expanded tool format: {type(expanded_tools[0])}, keys: {list(expanded_tools[0].keys()) if isinstance(expanded_tools[0], dict) else 'not a dict'}"
|
||||
)
|
||||
|
||||
# Update tools for filtering
|
||||
tools = expanded_tools
|
||||
original_tool_count = len(tools)
|
||||
|
|
@ -202,13 +201,10 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
# Always update tools and emit header (even if count unchanged)
|
||||
data["tools"] = filtered_tools
|
||||
|
||||
# Store filter stats for response header
|
||||
# Store filter stats for response header (following standard hook pattern)
|
||||
filter_stats = f"{original_tool_count}->{len(filtered_tools)}"
|
||||
|
||||
# Store in proxy_server_request (internal field not sent to providers)
|
||||
if "proxy_server_request" not in data:
|
||||
data["proxy_server_request"] = {}
|
||||
data["proxy_server_request"]["semantic_filter_stats"] = filter_stats
|
||||
_metadata_variable_name = self._get_metadata_variable_name(data)
|
||||
data[_metadata_variable_name]["litellm_semantic_filter_stats"] = filter_stats
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Semantic tool filter: {filter_stats} tools"
|
||||
|
|
@ -222,6 +218,24 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
)
|
||||
return None
|
||||
|
||||
async def async_post_call_response_headers_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
response: Any,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Add semantic filter stats to response headers.
|
||||
"""
|
||||
_metadata_variable_name = self._get_metadata_variable_name(data)
|
||||
filter_stats = data[_metadata_variable_name].get("litellm_semantic_filter_stats")
|
||||
if filter_stats:
|
||||
return {
|
||||
"x-litellm-semantic-filter": filter_stats
|
||||
}
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def initialize_from_config(
|
||||
config: Optional[Dict[str, Any]],
|
||||
|
|
@ -240,20 +254,16 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
|
||||
SemanticMCPToolFilter,
|
||||
)
|
||||
verbose_proxy_logger.info(f"🔍 initialize_from_config called: config={config}, llm_router={llm_router is not None}")
|
||||
|
||||
if not config or not config.get("enabled", False):
|
||||
verbose_proxy_logger.warning(f"❌ Semantic tool filter not enabled: config={config}")
|
||||
verbose_proxy_logger.debug("Semantic tool filter not enabled in config")
|
||||
return None
|
||||
|
||||
if llm_router is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"❌ Cannot initialize semantic filter: llm_router is None"
|
||||
"Cannot initialize semantic filter: llm_router is None"
|
||||
)
|
||||
return None
|
||||
|
||||
verbose_proxy_logger.info("✅ Config and router available, creating filter...")
|
||||
|
||||
try:
|
||||
|
||||
embedding_model = config.get(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue