Skip to content

feat: add opt-in ToolCache for caching tool invocation results in ToolInvoker and Agent #11824

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Kunal-Somani wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat: add opt-in ToolCache for caching tool invocation results in ToolInvoker and Agent #11824

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d761d0d
feat: add opt-in ToolCache for caching tool invocation results in Too…
Kunal-Somani Jun 29, 2026
88cc7de
fix: add missing SPDX license header to tool_cache.py
Kunal-Somani Jun 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 41 additions & 1 deletion haystack/components/agents/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
flatten_tools_or_toolsets,
serialize_tools_or_toolset,
)
from haystack.tools.tool_cache import ToolCache
from haystack.utils import _deserialize_value_with_schema
from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
from haystack.utils.deserialization import deserialize_component_inplace
Expand Down Expand Up @@ -206,6 +207,21 @@ def translate(
print(result["last_message"].text)
```

#### Using a `tool_cache` to avoid redundant tool calls

When an Agent loops over several steps, it can end up calling the same read-only tool with the same
arguments more than once (for example, re-fetching a document it already fetched a few steps earlier).
Pass a `ToolCache` to avoid re-invoking tools that have opted in via `Tool(..., cacheable=True)`:

```python
from haystack.components.agents import Agent
from haystack.tools.tool_cache import ToolCache

cache = ToolCache(ttl_seconds=3600, scope="agent_run")
agent = Agent(chat_generator=OpenAIChatGenerator(), tools=[my_cacheable_tool], tool_cache=cache)
result = agent.run(messages=[ChatMessage.from_user("...")])
# result["tool_cache_stats"] -> {"hits": ..., "misses": ..., "calls_saved": ...}
```
"""

def __init__(
Expand All @@ -223,6 +239,7 @@ def __init__(
raise_on_tool_invocation_failure: bool = False,
tool_invoker_kwargs: dict[str, Any] | None = None,
confirmation_strategies: dict[str | tuple[str, ...], ConfirmationStrategy] | None = None,
tool_cache: ToolCache | None = None,
) -> None:
"""
Initialize the agent component.
Expand Down Expand Up @@ -254,6 +271,12 @@ def __init__(
If set to False, the exception will be turned into a chat message and passed to the LLM.
:param tool_invoker_kwargs: Additional keyword arguments to pass to the ToolInvoker.
:param confirmation_strategies: A dictionary mapping tool names to ConfirmationStrategy instances.
:param tool_cache: Optional `ToolCache` instance shared across the Agent's run. Passed straight through
to the internal `ToolInvoker`, so tool calls for any `Tool` with `cacheable=True` are looked up
in the cache before invocation and stored after a successful invocation. When set, the dict returned
by `run()`/`run_async()` includes a `"tool_cache_stats"` key with the cache's hit/miss counts for
the run. Not included in `to_dict()`/`from_dict()` serialization, since cache backends are runtime
objects rather than serializable configuration — see `ToolInvoker.tool_cache` for the same caveat.
:raises TypeError: If the chat_generator does not support tools parameter in its run method.
:raises ValueError: If the exit_conditions are not valid.
:raises ValueError: If any `user_prompt` variable overlaps with the `state_schema` or `run` method parameters.
Expand Down Expand Up @@ -298,10 +321,13 @@ def __init__(
self.max_agent_steps = max_agent_steps
self.raise_on_tool_invocation_failure = raise_on_tool_invocation_failure
self.streaming_callback = streaming_callback
self.tool_cache = tool_cache

# Set input and output types for the component based on the State schema
self._run_method_params = _get_run_method_params(self)
output_types = {"last_message": ChatMessage}
output_types: dict[str, Any] = {"last_message": ChatMessage}
if self.tool_cache is not None:
output_types["tool_cache_stats"] = dict[str, int]
for param, config in self.state_schema.items():
output_types[param] = config["type"]
# Skip setting input types for parameters that are already in the run method
Expand All @@ -327,6 +353,7 @@ def __init__(
resolved_tool_invoker_kwargs = {
"tools": self.tools,
"raise_on_failure": self.raise_on_tool_invocation_failure,
"tool_cache": self.tool_cache,
**(tool_invoker_kwargs or {}),
}
self._tool_invoker = ToolInvoker(**resolved_tool_invoker_kwargs)
Expand Down Expand Up @@ -407,6 +434,11 @@ def to_dict(self) -> dict[str, Any]:
"""
Serialize the component to a dictionary.

Note: `tool_cache`, if set, is intentionally NOT included — cache backends are runtime objects
rather than serializable configuration. A deserialized `Agent` will have `tool_cache=None` even
if the original instance had one set; callers that need caching after deserialization should
re-attach a `ToolCache` explicitly.

:returns: Dictionary with serialized data.
"""
return default_to_dict(
Expand Down Expand Up @@ -768,6 +800,8 @@ def run( # noqa: PLR0915
A dictionary with the following keys:
- "messages": List of all messages exchanged during the agent's run.
- "last_message": The last message exchanged during the agent's run.
- "tool_cache_stats": Only present if a `tool_cache` was configured; a dict with `hits`, `misses`,
and `calls_saved` counts for this run.
- Any additional keys defined in the `state_schema`.
:raises BreakpointException: If an agent breakpoint is triggered.
"""
Expand Down Expand Up @@ -961,6 +995,8 @@ def run( # noqa: PLR0915
result = {**exe_context.state.data}
if msgs := result.get("messages"):
result["last_message"] = msgs[-1]
if self.tool_cache is not None:
result["tool_cache_stats"] = self.tool_cache.stats.to_dict()
return result

async def run_async( # noqa: PLR0915
Expand Down Expand Up @@ -1011,6 +1047,8 @@ async def run_async( # noqa: PLR0915
A dictionary with the following keys:
- "messages": List of all messages exchanged during the agent's run.
- "last_message": The last message exchanged during the agent's run.
- "tool_cache_stats": Only present if a `tool_cache` was configured; a dict with `hits`, `misses`,
and `calls_saved` counts for this run.
- Any additional keys defined in the `state_schema`.
:raises BreakpointException: If an agent breakpoint is triggered.
"""
Expand Down Expand Up @@ -1204,6 +1242,8 @@ async def run_async( # noqa: PLR0915
result = {**exe_context.state.data}
if msgs := result.get("messages"):
result["last_message"] = msgs[-1]
if self.tool_cache is not None:
result["tool_cache_stats"] = self.tool_cache.stats.to_dict()
return result

def _check_exit_conditions(self, llm_messages: list[ChatMessage], tool_messages: list[ChatMessage]) -> bool:
Expand Down
119 changes: 108 additions & 11 deletions haystack/components/tools/tool_invoker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
)
from haystack.tools.errors import ToolInvocationError
from haystack.tools.parameters_schema_utils import _unwrap_optional
from haystack.tools.tool_cache import ToolCache
from haystack.tracing.utils import _serializable_value
from haystack.utils.callable_serialization import deserialize_callable, serialize_callable

Expand Down Expand Up @@ -179,6 +180,18 @@ def dummy_weather_function(city: str):
result = invoker.run(messages=[message])

print(result)
```

Usage example with a ToolCache (opt-in per-tool caching, see `Tool.cacheable`):
```python
from haystack.tools.tool_cache import ToolCache

cache = ToolCache(ttl_seconds=3600, scope="agent_run")
cacheable_tool = Tool(..., cacheable=True)
invoker = ToolInvoker(tools=[cacheable_tool], tool_cache=cache)
# Identical calls to `cacheable_tool` within the cache's TTL are served from cache;
# `cache.stats` tracks hits/misses for inspection after the run.
```
"""

def __init__(
Expand All @@ -190,6 +203,7 @@ def __init__(
*,
enable_streaming_callback_passthrough: bool = False,
max_workers: int = 4,
tool_cache: ToolCache | None = None,
) -> None:
"""
Initialize the ToolInvoker component.
Expand All @@ -216,6 +230,12 @@ def __init__(
:param max_workers:
The maximum number of workers to use in the thread pool executor.
This also decides the maximum number of concurrent tool invocations.
:param tool_cache:
Optional `ToolCache` instance. When provided, tool calls for any `Tool` with `cacheable=True`
are looked up in the cache before invocation and the result is stored in the cache after a
successful invocation. Tools with `cacheable=False` (the default) always invoke normally,
regardless of whether a `tool_cache` is configured. Not included in `to_dict()`/`from_dict()`
serialization, since cache backends are runtime objects rather than serializable configuration.
:raises ValueError:
If no tools are provided or if duplicate tool names are found.
"""
Expand All @@ -225,6 +245,7 @@ def __init__(
self.max_workers = max_workers
self.raise_on_failure = raise_on_failure
self.convert_result_to_json_string = convert_result_to_json_string
self.tool_cache = tool_cache

self._tools_with_names = self._validate_and_prepare_tools(tools)
self._is_warmed_up = False
Expand Down Expand Up @@ -252,6 +273,21 @@ def _runner() -> Any:

return _runner

@staticmethod
def _make_cached_result_invoke(cached_result: Any) -> Callable[[], Any]:
"""
Create a zero-arg callable that simply returns an already-cached result.

Used in place of `_make_context_bound_invoke` on a cache hit, so the surrounding
ThreadPoolExecutor/future-based control flow in `run`/`run_async` stays uniform
regardless of whether a given call was actually invoked or served from cache.
"""

def _runner() -> Any:
return cached_result

return _runner

@staticmethod
def _validate_and_prepare_tools(tools: ToolsType) -> dict[str, Tool]:
"""
Expand Down Expand Up @@ -494,6 +530,34 @@ def _create_tool_result_streaming_chunk(tool_messages: list[ChatMessage], tool_c
meta={"tool_result": tool_messages[-1].tool_call_results[0].result, "tool_call": tool_call},
)

def _check_cache(self, tool_to_invoke: Tool, tool_call: ToolCall) -> tuple[bool, Any]:
"""
Look up a tool call in the configured ToolCache, if any.

Only consulted when `self.tool_cache` is set AND the tool itself opted in via `Tool.cacheable=True`.
A tool with `cacheable=False` (the default) is never looked up, regardless of whether a cache is configured.

:param tool_to_invoke: The Tool that would be invoked.
:param tool_call: The ToolCall containing the arguments to look up.
:returns: A tuple `(hit, cached_result)`. `hit` is always False if no cache is configured or the tool
is not cacheable.
"""
if self.tool_cache is None or not getattr(tool_to_invoke, "cacheable", False):
return False, None
return self.tool_cache.get(tool_to_invoke.name, tool_call.arguments)

def _store_in_cache(self, tool_to_invoke: Tool, tool_call: ToolCall, result: Any) -> None:
"""
Store a successful tool invocation result in the configured ToolCache, if any.

:param tool_to_invoke: The Tool that was invoked.
:param tool_call: The ToolCall containing the arguments used.
:param result: The result to store.
"""
if self.tool_cache is None or not getattr(tool_to_invoke, "cacheable", False):
return
self.tool_cache.set(tool_to_invoke.name, tool_call.arguments, result)

def _prepare_tool_call_params(
self,
*,
Expand All @@ -506,6 +570,10 @@ def _prepare_tool_call_params(
"""
Prepare tool call parameters for execution and collect any error messages.

Also performs the cache lookup for cacheable tools: each entry in the returned `tool_call_params`
list carries `cache_hit` (bool) and `cached_result` (Any), so the caller (`run`/`run_async`) can
skip actual invocation on a hit while keeping a uniform future-based execution path.

:param messages_with_tool_calls: Messages containing tool calls to process
:param state: The current state for argument injection
:param streaming_callback: Optional streaming callback to inject
Expand Down Expand Up @@ -544,7 +612,19 @@ def _prepare_tool_call_params(
):
final_args["streaming_callback"] = streaming_callback

tool_call_params.append({"tool_to_invoke": tool_to_invoke, "final_args": final_args})
# Cache lookup happens against the original LLM-provided arguments (tool_call.arguments),
# not final_args, so that state-injected/streaming_callback-injected values (which can vary
# run-to-run without changing the logical call) don't accidentally fragment the cache key.
cache_hit, cached_result = self._check_cache(tool_to_invoke, tool_call)

tool_call_params.append(
{
"tool_to_invoke": tool_to_invoke,
"final_args": final_args,
"cache_hit": cache_hit,
"cached_result": cached_result,
}
)
tool_calls.append(tool_call)

return tool_calls, tool_call_params, error_messages
Expand Down Expand Up @@ -645,15 +725,19 @@ def run(
if not tool_call_params:
return {"tool_messages": tool_messages, "state": state}

# 2) Execute valid tool calls in parallel
# 2) Execute valid tool calls in parallel (cache hits get a trivial pass-through callable,
# so the executor/future machinery stays identical regardless of cache outcome)
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
futures = []
for params in tool_call_params:
callable_ = self._make_context_bound_invoke(params["tool_to_invoke"], params["final_args"])
if params["cache_hit"]:
callable_ = self._make_cached_result_invoke(params["cached_result"])
else:
callable_ = self._make_context_bound_invoke(params["tool_to_invoke"], params["final_args"])
futures.append(executor.submit(callable_))

# 3) Gather and process results: handle errors and merge outputs into state
for future, tool_call in zip(futures, tool_calls, strict=True):
# 3) Gather and process results: handle errors, store cache misses, and merge outputs into state
for future, tool_call, params in zip(futures, tool_calls, tool_call_params, strict=True):
result = future.result()

if isinstance(result, ToolInvocationError):
Expand All @@ -663,9 +747,11 @@ def run(
logger.error("{error_exception}", error_exception=result)
tool_messages.append(ChatMessage.from_tool(tool_result=str(result), origin=tool_call, error=True))
else:
# b) In case of success, merge outputs into state
# b) In case of success, store cache misses, merge outputs into state
try:
tool_to_invoke = tools_with_names[tool_call.tool_name]
if not params["cache_hit"]:
self._store_in_cache(tool_to_invoke, tool_call, result)
self._merge_tool_outputs(tool=tool_to_invoke, result=result, state=state)
tool_messages.append(
self._prepare_tool_result_message(
Expand Down Expand Up @@ -781,17 +867,21 @@ async def run_async(
if not tool_call_params:
return {"tool_messages": tool_messages, "state": state}

# 2) Execute valid tool calls in parallel
# 2) Execute valid tool calls in parallel (cache hits get a trivial pass-through callable,
# so the executor/future machinery stays identical regardless of cache outcome)
tool_call_tasks = []
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
for params in tool_call_params:
loop = asyncio.get_running_loop()
callable_ = ToolInvoker._make_context_bound_invoke(params["tool_to_invoke"], params["final_args"])
if params["cache_hit"]:
callable_ = ToolInvoker._make_cached_result_invoke(params["cached_result"])
else:
callable_ = ToolInvoker._make_context_bound_invoke(params["tool_to_invoke"], params["final_args"])
tool_call_tasks.append(loop.run_in_executor(executor, callable_))

# 3) Gather and process results: handle errors and merge outputs into state
# 3) Gather and process results: handle errors, store cache misses, and merge outputs into state
tool_results = await asyncio.gather(*tool_call_tasks)
for tool_result, tool_call in zip(tool_results, tool_calls, strict=True):
for tool_result, tool_call, params in zip(tool_results, tool_calls, tool_call_params, strict=True):
# a) This is an error, create error Tool message
if isinstance(tool_result, ToolInvocationError):
if self.raise_on_failure:
Expand All @@ -801,9 +891,11 @@ async def run_async(
ChatMessage.from_tool(tool_result=str(tool_result), origin=tool_call, error=True)
)
else:
# b) In case of success, merge outputs into state
# b) In case of success, store cache misses, merge outputs into state
try:
tool_to_invoke = tools_with_names[tool_call.tool_name]
if not params["cache_hit"]:
self._store_in_cache(tool_to_invoke, tool_call, tool_result)
self._merge_tool_outputs(tool=tool_to_invoke, result=tool_result, state=state)
tool_messages.append(
self._prepare_tool_result_message(
Expand Down Expand Up @@ -839,6 +931,11 @@ def to_dict(self) -> dict[str, Any]:
"""
Serializes the component to a dictionary.

Note: `tool_cache`, if set, is intentionally NOT included — cache backends are runtime objects
(in-memory state, possibly external connections) rather than serializable configuration. A
deserialized `ToolInvoker` will have `tool_cache=None` even if the original instance had one set;
callers that need caching after deserialization should re-attach a `ToolCache` explicitly.

:returns:
Dictionary with serialized data.
"""
Expand Down
Loading
Loading