add mcp filters to client
This commit is contained in:
@@ -7,7 +7,7 @@
|
||||
"""MCP (Model Context Protocol) client for integrating external tools with LLMs."""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, TypeAlias
|
||||
from typing import Any, Dict, List, TypeAlias, Callable
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@@ -46,17 +46,24 @@ class MCPClient(BaseObject):
|
||||
def __init__(
|
||||
self,
|
||||
server_params: ServerParameters,
|
||||
tools_filter: List[str] | None = None,
|
||||
tools_output_filters: Dict[str, Callable[[Any], Any]] | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the MCP client with server parameters.
|
||||
|
||||
Args:
|
||||
server_params: Server connection parameters (stdio or SSE).
|
||||
tools_filter: Optional list of tool names to register. If None, all tools are registered.
|
||||
tools_output_filters: Optional dict mapping tool names to filter functions that process tool outputs.
|
||||
Each filter function receives the raw tool output (any type) and returns the processed output (any type).
|
||||
**kwargs: Additional arguments passed to the parent BaseObject.
|
||||
"""
|
||||
super().__init__(**kwargs)
|
||||
self._server_params = server_params
|
||||
self._session = ClientSession
|
||||
self._tools_filter = tools_filter
|
||||
self._tools_output_filters = tools_output_filters or {}
|
||||
|
||||
if isinstance(server_params, StdioServerParameters):
|
||||
self._client = stdio_client
|
||||
@@ -264,13 +271,24 @@ class MCPClient(BaseObject):
|
||||
else:
|
||||
# logger.debug(f"Non-text result content: '{content}'")
|
||||
pass
|
||||
logger.info(f"Tool '{function_name}' completed successfully")
|
||||
logger.debug(f"Final response: {response}")
|
||||
else:
|
||||
logger.error(f"Error getting content from {function_name} results.")
|
||||
|
||||
final_response = response if len(response) else "Sorry, could not call the mcp tool"
|
||||
await result_callback(final_response)
|
||||
# Apply output filter if configured for this tool
|
||||
if function_name in self._tools_output_filters:
|
||||
try:
|
||||
response = self._tools_output_filters[function_name](response)
|
||||
logger.debug(f"Final response (after filter): {response}")
|
||||
|
||||
except Exception:
|
||||
logger.error("Error applying output filter")
|
||||
response = ""
|
||||
|
||||
if response and isinstance(response, str):
|
||||
logger.info(f"Tool '{function_name}' completed successfully")
|
||||
logger.debug(f"Final response: {response}")
|
||||
|
||||
await result_callback(response)
|
||||
|
||||
async def _list_tools_helper(self, session):
|
||||
available_tools = await session.list_tools()
|
||||
@@ -283,6 +301,12 @@ class MCPClient(BaseObject):
|
||||
|
||||
for tool in available_tools.tools:
|
||||
tool_name = tool.name
|
||||
|
||||
# Apply tools filter if configured
|
||||
if self._tools_filter and tool_name not in self._tools_filter:
|
||||
logger.debug(f"Skipping tool '{tool_name}' - not in allowed tools list")
|
||||
continue
|
||||
|
||||
logger.debug(f"Processing tool: {tool_name}")
|
||||
logger.debug(f"Tool description: {tool.description}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user