feat: add async_execute_tool_calls hook to CustomLogger

Adds a simplified tool execution hook pattern to the LiteLLM framework:

- ToolCallResult NamedTuple in custom_logger.py for structured tool results
- async_execute_tool_calls() method on CustomLogger base class
- _call_tool_execution_hooks() in BaseLLMHTTPHandler that aggregates
  results from all callbacks, constructs follow-up messages (preserving
  thinking blocks), and makes the follow-up API request
- _filter_handled_tool_calls() to allow multiple callbacks to handle
  different tool calls from the same response

This is a framework-level change that moves agentic loop orchestration
into the HTTP handler, so individual CustomLogger implementations only
need to execute tool calls and return results.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Michael Cusack 2026-03-04 19:49:13 -08:00
parent 3c389ad6f7
commit 72ddc18757
2 changed files with 308 additions and 1 deletions

View file

@ -8,6 +8,7 @@ from typing import (
AsyncGenerator,
Dict,
List,
NamedTuple,
Optional,
Tuple,
Union,
@ -43,6 +44,9 @@ if TYPE_CHECKING:
MCPPreCallRequestObject,
MCPPreCallResponseObject,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.router import PreRoutingHookResponse
Span = Union[_Span, Any]
@ -56,6 +60,7 @@ else:
MCPDuringCallRequestObject = Any
MCPDuringCallResponseObject = Any
PreRoutingHookResponse = Any
AnthropicMessagesResponse = Any
_BASE64_INLINE_PATTERN = re.compile(
@ -64,6 +69,19 @@ _BASE64_INLINE_PATTERN = re.compile(
)
class ToolCallResult(NamedTuple):
"""Result of executing a single tool call via async_execute_tool_calls."""
tool_call_id: str
"""The id of the tool_use block that was executed."""
content: str
"""Text result to return to the model."""
is_error: bool
"""Whether this result represents an error."""
class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class
# Class variables or attributes
def __init__(
@ -533,6 +551,34 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
return None
#########################################################
# TOOL EXECUTION HOOKS (simplified tool interception)
#########################################################
async def async_execute_tool_calls(
self,
response: Union["AnthropicMessagesResponse", ModelResponse],
kwargs: Dict,
) -> List[ToolCallResult]:
"""
Detect and execute tool calls in the model response.
This is the simplified alternative to the two-step
async_should_run_agentic_loop / async_run_agentic_loop pattern.
Callbacks only need to detect tool calls and return results — the
framework handles message construction, thinking block preservation,
max_tokens adjustment, kwargs cleanup, and follow-up API requests.
Args:
response: Model response (AnthropicMessagesResponse dict, or ModelResponse)
kwargs: Full request kwargs (includes custom_llm_provider, tools, etc.)
Returns:
List of ToolCallResult for tool calls this callback handled.
Empty list means nothing was handled (skip this callback).
"""
return []
#########################################################
# AGENTIC LOOP HOOKS (for litellm.messages + future completion support)
#########################################################

View file

@ -1984,7 +1984,22 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
)
# Call agentic completion hooks
# Call simplified tool execution hooks first (new pattern)
tool_exec_response = await self._call_tool_execution_hooks(
response=initial_response,
model=model,
messages=messages,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream or False,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
if tool_exec_response is not None:
return tool_exec_response
# Call legacy agentic completion hooks (old two-step pattern)
final_response = await self._call_agentic_completion_hooks(
response=initial_response,
model=model,
@ -4390,6 +4405,252 @@ class BaseLLMHTTPHandler:
return stream, data
return stream, data
async def _call_tool_execution_hooks(
self,
response: Any,
model: str,
messages: List[Dict],
anthropic_messages_provider_config: "BaseAnthropicMessagesConfig",
anthropic_messages_optional_request_params: Dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
custom_llm_provider: str,
kwargs: Dict,
) -> Optional[Any]:
"""
Check all callbacks for tool calls to execute via the simplified
async_execute_tool_calls hook.
Aggregates results from all callbacks, then handles the agentic loop.
Multiple callbacks can handle different tool calls from the same response.
Returns the final response after tool execution, or None if no hooks fired.
"""
from litellm.integrations.custom_logger import CustomLogger, ToolCallResult
callbacks = litellm.callbacks + (
logging_obj.dynamic_success_callbacks or []
)
all_results: List[ToolCallResult] = []
working_response = response # progressively filtered
# Ensure custom_llm_provider is in kwargs for callbacks
kwargs_with_provider = kwargs.copy() if kwargs else {}
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
for callback in callbacks:
try:
if not isinstance(callback, CustomLogger):
continue
results = await callback.async_execute_tool_calls(
response=working_response,
kwargs=kwargs_with_provider,
)
if results:
all_results.extend(results)
# Remove handled tool_use blocks so next callback only sees unhandled ones
handled_ids = {r.tool_call_id for r in results}
working_response = self._filter_handled_tool_calls(
working_response, handled_ids
)
except Exception as e:
verbose_logger.exception(
f"LiteLLM.ToolExecutionHookError: Exception in tool execution hooks: {str(e)}"
)
if not all_results:
return None
# Framework builds messages and makes follow-up request
return await self._complete_tool_execution_loop(
response=response, # original unfiltered response
results=all_results,
model=model,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs,
)
@staticmethod
def _filter_handled_tool_calls(
response: Any,
handled_ids: set,
) -> Any:
"""
Return a copy of ``response`` with tool_use blocks whose ids are in
``handled_ids`` removed. This lets subsequent callbacks only see
unhandled tool calls.
"""
if isinstance(response, dict):
content = response.get("content")
if content is None:
return response
filtered = [
b for b in content
if not (
(isinstance(b, dict) and b.get("type") == "tool_use" and b.get("id") in handled_ids)
or (hasattr(b, "type") and getattr(b, "type", None) == "tool_use" and getattr(b, "id", None) in handled_ids)
)
]
return {**response, "content": filtered}
# Object-style response
if hasattr(response, "content") and response.content is not None:
import copy
filtered = [
b for b in response.content
if not (
getattr(b, "type", None) == "tool_use"
and getattr(b, "id", None) in handled_ids
)
]
new_resp = copy.copy(response)
new_resp.content = filtered
return new_resp
return response
async def _complete_tool_execution_loop(
self,
response: Any,
results: List,
model: str,
messages: List[Dict],
anthropic_messages_optional_request_params: Dict,
logging_obj: "LiteLLMLoggingObj",
kwargs: Dict,
) -> Any:
"""
Framework handles ALL the plumbing after callbacks return ToolCallResults:
1. Extract thinking blocks + tool_use blocks from original response
2. Build assistant message (thinking + tool_use)
3. Build tool_result user message (matching results by tool_call_id)
4. Adjust max_tokens for thinking token usage
5. Clean kwargs of internal keys
6. Make follow-up API request
7. Return final response
"""
from litellm.anthropic_interface import messages as anthropic_messages
from litellm.litellm_core_utils.core_helpers import filter_internal_params
# ---- 1. Extract thinking blocks and tool_use blocks from response ----
if isinstance(response, dict):
content = response.get("content", [])
else:
content = getattr(response, "content", None) or []
thinking_blocks: List[Dict] = []
tool_use_blocks: List[Dict] = []
handled_ids = {r.tool_call_id for r in results}
for block in content:
if isinstance(block, dict):
btype = block.get("type")
bid = block.get("id")
else:
btype = getattr(block, "type", None)
bid = getattr(block, "id", None)
if btype in ("thinking", "redacted_thinking"):
if isinstance(block, dict):
thinking_blocks.append(block)
else:
normalized: Dict[str, Any] = {"type": btype}
for attr in ("thinking", "data", "signature"):
if hasattr(block, attr):
normalized[attr] = getattr(block, attr)
thinking_blocks.append(normalized)
elif btype == "tool_use" and bid in handled_ids:
if isinstance(block, dict):
tool_use_blocks.append(block)
else:
tool_use_blocks.append({
"type": "tool_use",
"id": bid,
"name": getattr(block, "name", ""),
"input": getattr(block, "input", {}),
})
# ---- 2. Build assistant message (thinking first, then tool_use) ----
assistant_content: List[Dict] = []
if thinking_blocks:
assistant_content.extend(thinking_blocks)
assistant_content.extend(tool_use_blocks)
assistant_message = {"role": "assistant", "content": assistant_content}
# ---- 3. Build tool_result user message ----
results_by_id = {r.tool_call_id: r for r in results}
tool_result_blocks = []
for tu in tool_use_blocks:
tc_id = tu["id"]
r = results_by_id.get(tc_id)
block: Dict[str, Any] = {
"type": "tool_result",
"tool_use_id": tc_id,
"content": r.content if r else "",
}
if r and r.is_error:
block["is_error"] = True
tool_result_blocks.append(block)
user_message = {"role": "user", "content": tool_result_blocks}
# ---- 4. Prepare max_tokens (subtract thinking usage if present) ----
max_tokens = anthropic_messages_optional_request_params.get(
"max_tokens",
kwargs.get("max_tokens", 1024),
)
# If thinking is enabled, subtract thinking token usage from max_tokens
if isinstance(response, dict):
usage = response.get("usage", {})
else:
usage = getattr(response, "usage", None)
if usage and not isinstance(usage, dict):
usage = getattr(usage, "__dict__", {})
if usage:
# cache_creation_input_tokens is used in some responses; thinking tokens
# are not currently reported separately, but this is where we'd adjust.
pass
# ---- 5. Clean kwargs for follow-up request ----
kwargs_for_followup = filter_internal_params(kwargs)
# ---- 6. Resolve full model name ----
full_model_name = model
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
full_model_name = agentic_params.get("model", model)
optional_params_without_max_tokens = {
k: v
for k, v in anthropic_messages_optional_request_params.items()
if k != "max_tokens"
}
follow_up_messages = messages + [assistant_message, user_message]
verbose_logger.debug(
f"ToolExecutionLoop: Making follow-up request with "
f"{len(results)} tool result(s), model={full_model_name}"
)
# ---- 7. Make follow-up API request ----
final_response = await anthropic_messages.acreate(
max_tokens=max_tokens,
messages=follow_up_messages,
model=full_model_name,
**optional_params_without_max_tokens,
**kwargs_for_followup,
)
return final_response
async def _call_agentic_completion_hooks(
self,
response: Any,