Add MCPTool/MCPToolset warm_up #2384
Conversation
github-actions
bot
added
integration:mcp
type:documentation
|
A few notes: eager_connect defaults to False so no-code platforms don't have to specify each mcp tool/toolset as eager_connect false. With Tool.warm_up()/Toolset.warm_up() in place in Haystack 2.19 users that prefer eager tool init can still opt in explicitly and long-term default being False makes more sense. (I'll add this explicitly in release notes) I used threading.RLock() as _ensure_connected() mutates important shared state (_client, _worker, schema) and can be entered from synchronous invoke, warm_up, and async call dispatched from different threads or loops. This re-entrant lock is cheap insurance against double-connect races, leaked session managers and incorrect tool parameters fetched and overwritten. |
vblagoje
requested review from
mpangrazzi and
sjrl
and removed request for
a team
|
@mpangrazzi I also added @sjrl for review because he has more context about the whole Tool(set) warm_up initiative. I also want to update examples to showcase more use of Agent and warm_up but I left that change for separate PR for easier Pr digest - I have the changes locally and will open a separate PR for examples. |
| # Connect on first use if eager_connect is turned off | ||
| if not self._eager_connect: | ||
| self._ensure_connected() |
There was a problem hiding this comment.
We don't need this right since we would expect a user to connect via warm_up?
There was a problem hiding this comment.
At the very least if we do want this behavior I think we should just do
| # Connect on first use if eager_connect is turned off | |
| if not self._eager_connect: | |
| self._ensure_connected() | |
| # Connect on first use if eager_connect is turned off | |
| self.warm_up() |
and let warm_up handle this if logic
| if not self._eager_connect: | ||
| self._ensure_connected() |
There was a problem hiding this comment.
Same here if we want this behavior just call warm_up
| if not self._eager_connect: | |
| self._ensure_connected() | |
| self.warm_up() |
| return | ||
| self._ensure_connected() | ||
|
|
||
| def _ensure_connected(self) -> None: |
There was a problem hiding this comment.
If we end up just calling warm_up like I suggest in my above comments let's remove this function _ensure_connected and just directly put all its logic into warm_up
There was a problem hiding this comment.
Ok yes, i think it was a remnant of previous work where I also wanted to reconnect when broken connection was detected. Let me double check, thanks for raising it
| if self._client is None: | ||
| raise MCPConnectionError(message="Not connected to an MCP server", operation="call_tool") |
There was a problem hiding this comment.
Would it be better to use an ignore or an assert statement since we know this error can never be raised?
| try: | ||
| return await asyncio.wait_for(self._client.call_tool(self.name, kwargs), timeout=self._invocation_timeout) | ||
| self.warm_up() | ||
| client = cast(MCPClient, self._client) |
There was a problem hiding this comment.
Seems like you are using a different approach here to get around the mypy error than here. Let's make it consistent.
| if self._eager_connect: | ||
| return |
There was a problem hiding this comment.
Not a big deal but this isn't needed really since the check below
if self._client is not None:
returnwill also just return early right?
| with self._lock: | ||
| if self._client is not None: | ||
| return | ||
| client = self._server_info.create_client() | ||
| worker = _MCPClientSessionManager(client, timeout=self._connection_timeout) | ||
| tools = worker.tools() | ||
| tool = next((t for t in tools if t.name == self.name), None) | ||
| if tool is None: | ||
| available = [t.name for t in tools] | ||
| msg = f"Tool '{self.name}' not found on server. Available tools: {', '.join(available)}" | ||
| raise MCPToolNotFoundError(msg, tool_name=self.name, available_tools=available) | ||
| # Publish connection and tighten parameters for better prompting | ||
| self._client = client | ||
| self._worker = worker | ||
| try: | ||
| self.parameters = tool.inputSchema | ||
| except Exception as e: | ||
| logger.debug(f"TOOL: Could not update strict parameters after connect: {e!s}") |
There was a problem hiding this comment.
Hey I'm noticing that similar code to this lives in the init method starting at current line 855. Could we make a shared method that get's re-used in both places rather than having duplicate code?
There was a problem hiding this comment.
If possible, one way could be to just call self.warm_up in the init method if self._eager_connect is True
| self.connection_timeout = connection_timeout | ||
| self.invocation_timeout = invocation_timeout | ||
| self.eager_connect = eager_connect | ||
| self.warmup_called = False |
There was a problem hiding this comment.
I think I'd make the warmup_called variable private
| Call this before handing the toolset to components like ``ToolInvoker`` so that | ||
| each tool's schema is available without performing a real invocation. |
There was a problem hiding this comment.
This is a bit misleading since users could also:
- Call
ToolInvoker.warm_up()instead ofMCPToolset.warm_up() - Call
Pipeline.warm_up()if theToolInvokeris part of a pipeline
| if self.eager_connect or self.warmup_called: | ||
| return |
There was a problem hiding this comment.
If we make the change in this comment https://github.com/deepset-ai/haystack-core-integrations/pull/2384/files#r2451228527 then we can update this to
| if self.eager_connect or self.warmup_called: | |
| return | |
| if self.warmup_called: | |
| return |
|
Great feedback @sjrl - I also tested this latest commit against itinerary agent |
Why:
Separates pipeline validation from MCP server connection. Allows pipelines to validate successfully without requiring server connectivity, then connect during pipeline warm_up.
Keep in
draft stage until deepset-ai/haystack#9856 is merge and Haystack 2.19 released.What:
eager_connectparameter toMCPToolandMCPToolset(defaults toFalse)_ensure_connected()with thread-safe lazy initializationwarm_up()methods for connecting during pipeline warm_up phaseHow can it be used:
How did you test it:
Notes for the reviewer:
warm_up()is idempotent